From 751e250f70cf446ae342c8a860d92f6a8b78261a Mon Sep 17 00:00:00 2001 From: Samarth Agrawal <41808786+SammyAgrawal@users.noreply.github.com> Date: Mon, 15 Sep 2025 23:48:48 -0700 Subject: [PATCH 001/123] fixed bug in defining embed dim for UNet1D (#12111) * fixed bug in defining embed dim * matched 1d temb process to 2d * Update src/diffusers/models/unets/unet_1d.py Co-authored-by: Dhruv Nair --------- Co-authored-by: Dhruv Nair --- src/diffusers/models/unets/unet_1d.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/src/diffusers/models/unets/unet_1d.py b/src/diffusers/models/unets/unet_1d.py index 4f57f3349b..4c4c528a59 100644 --- a/src/diffusers/models/unets/unet_1d.py +++ b/src/diffusers/models/unets/unet_1d.py @@ -82,6 +82,7 @@ class UNet1DModel(ModelMixin, ConfigMixin): out_channels: int = 2, extra_in_channels: int = 0, time_embedding_type: str = "fourier", + time_embedding_dim: Optional[int] = None, flip_sin_to_cos: bool = True, use_timestep_embedding: bool = False, freq_shift: float = 0.0, @@ -100,15 +101,23 @@ class UNet1DModel(ModelMixin, ConfigMixin): # time if time_embedding_type == "fourier": + time_embed_dim = time_embedding_dim or block_out_channels[0] * 2 + if time_embed_dim % 2 != 0: + raise ValueError(f"`time_embed_dim` should be divisible by 2, but is {time_embed_dim}.") self.time_proj = GaussianFourierProjection( - embedding_size=8, set_W_to_weight=False, log=False, flip_sin_to_cos=flip_sin_to_cos + embedding_size=time_embed_dim // 2, set_W_to_weight=False, log=False, flip_sin_to_cos=flip_sin_to_cos ) - timestep_input_dim = 2 * block_out_channels[0] + timestep_input_dim = time_embed_dim elif time_embedding_type == "positional": + time_embed_dim = time_embedding_dim or block_out_channels[0] * 4 self.time_proj = Timesteps( block_out_channels[0], flip_sin_to_cos=flip_sin_to_cos, downscale_freq_shift=freq_shift ) timestep_input_dim = block_out_channels[0] + else: + raise ValueError( + f"{time_embedding_type} does not exist. Please make sure to use one of `fourier` or `positional`." + ) if use_timestep_embedding: time_embed_dim = block_out_channels[0] * 4 From 8c72cd12ee65e420c86a0724f0182f966f339a7e Mon Sep 17 00:00:00 2001 From: Sari Hleihil <57064636+sarihl@users.noreply.github.com> Date: Wed, 17 Sep 2025 02:41:05 +0300 Subject: [PATCH 002/123] Added LucyEditPipeline (#12340) * Added LucyEditPipeline * add import & stype missing copied from * Fix example doc string --------- Co-authored-by: yiyixuxu --- src/diffusers/__init__.py | 2 + src/diffusers/pipelines/__init__.py | 2 + src/diffusers/pipelines/lucy/__init__.py | 47 ++ .../pipelines/lucy/pipeline_lucy_edit.py | 735 ++++++++++++++++++ .../pipelines/lucy/pipeline_output.py | 20 + .../dummy_torch_and_transformers_objects.py | 15 + 6 files changed, 821 insertions(+) create mode 100644 src/diffusers/pipelines/lucy/__init__.py create mode 100644 src/diffusers/pipelines/lucy/pipeline_lucy_edit.py create mode 100644 src/diffusers/pipelines/lucy/pipeline_output.py diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index d96acc3818..167d39c6e8 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -495,6 +495,7 @@ else: "LTXImageToVideoPipeline", "LTXLatentUpsamplePipeline", "LTXPipeline", + "LucyEditPipeline", "Lumina2Pipeline", "Lumina2Text2ImgPipeline", "LuminaPipeline", @@ -1149,6 +1150,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: LTXImageToVideoPipeline, LTXLatentUpsamplePipeline, LTXPipeline, + LucyEditPipeline, Lumina2Pipeline, Lumina2Text2ImgPipeline, LuminaPipeline, diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index 8ed07a72e3..17f3fc909e 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -285,6 +285,7 @@ else: ] _import_structure["lumina"] = ["LuminaPipeline", "LuminaText2ImgPipeline"] _import_structure["lumina2"] = ["Lumina2Pipeline", "Lumina2Text2ImgPipeline"] + _import_structure["lucy"] = ["LucyEditPipeline"] _import_structure["marigold"].extend( [ "MarigoldDepthPipeline", @@ -682,6 +683,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: LEditsPPPipelineStableDiffusionXL, ) from .ltx import LTXConditionPipeline, LTXImageToVideoPipeline, LTXLatentUpsamplePipeline, LTXPipeline + from .lucy import LucyEditPipeline from .lumina import LuminaPipeline, LuminaText2ImgPipeline from .lumina2 import Lumina2Pipeline, Lumina2Text2ImgPipeline from .marigold import ( diff --git a/src/diffusers/pipelines/lucy/__init__.py b/src/diffusers/pipelines/lucy/__init__.py new file mode 100644 index 0000000000..580e1f37f3 --- /dev/null +++ b/src/diffusers/pipelines/lucy/__init__.py @@ -0,0 +1,47 @@ +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects # noqa F403 + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["pipeline_lucy_edit"] = ["LucyEditPipeline"] +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() + + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * + else: + from .pipeline_lucy_edit import LucyEditPipeline + +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) + + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/lucy/pipeline_lucy_edit.py b/src/diffusers/pipelines/lucy/pipeline_lucy_edit.py new file mode 100644 index 0000000000..69f69d5768 --- /dev/null +++ b/src/diffusers/pipelines/lucy/pipeline_lucy_edit.py @@ -0,0 +1,735 @@ +# Copyright 2025 The Wan Team and The HuggingFace Team. All rights reserved. +# Copyright 2025 The Decart AI Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Modifications by Decart AI Team: +# - Based on pipeline_wan.py, but with supports recieving a condition video appended to the channel dimension. + +import html +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +import regex as re +import torch +from PIL import Image +from transformers import AutoTokenizer, UMT5EncoderModel + +from ...callbacks import MultiPipelineCallbacks, PipelineCallback +from ...loaders import WanLoraLoaderMixin +from ...models import AutoencoderKLWan, WanTransformer3DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import is_ftfy_available, is_torch_xla_available, logging, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ...video_processor import VideoProcessor +from ..pipeline_utils import DiffusionPipeline +from .pipeline_output import LucyPipelineOutput + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +if is_ftfy_available(): + import ftfy + + +EXAMPLE_DOC_STRING = """ + Examples: + ```python + >>> from typing import List + + >>> import torch + >>> from PIL import Image + + >>> from diffusers import AutoencoderKLWan, LucyEditPipeline + >>> from diffusers.utils import export_to_video, load_video + + >>> # Arguments + >>> url = "https://d2drjpuinn46lb.cloudfront.net/painter_original_edit.mp4" + >>> prompt = "Change the apron and blouse to a classic clown costume: satin polka-dot jumpsuit in bright primary colors, ruffled white collar, oversized pom-pom buttons, white gloves, oversized red shoes, red foam nose; soft window light from left, eye-level medium shot, natural folds and fabric highlights." + >>> negative_prompt = "" + >>> num_frames = 81 + >>> height = 480 + >>> width = 832 + + + >>> # Load video + >>> def convert_video(video: List[Image.Image]) -> List[Image.Image]: + ... video = load_video(url)[:num_frames] + ... video = [video[i].resize((width, height)) for i in range(num_frames)] + ... return video + + + >>> video = load_video(url, convert_method=convert_video) + + >>> # Load model + >>> model_id = "decart-ai/Lucy-Edit-Dev" + >>> vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32) + >>> pipe = LucyEditPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16) + >>> pipe.to("cuda") + + >>> # Generate video + >>> output = pipe( + ... prompt=prompt, + ... video=video, + ... negative_prompt=negative_prompt, + ... height=480, + ... width=832, + ... num_frames=81, + ... guidance_scale=5.0, + ... ).frames[0] + + >>> # Export video + >>> export_to_video(output, "output.mp4", fps=24) + ``` +""" + + +def basic_clean(text): + text = ftfy.fix_text(text) + text = html.unescape(html.unescape(text)) + return text.strip() + + +def whitespace_clean(text): + text = re.sub(r"\s+", " ", text) + text = text.strip() + return text + + +def prompt_clean(text): + text = whitespace_clean(basic_clean(text)) + return text + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents +def retrieve_latents( + encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" +): + if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": + return encoder_output.latent_dist.sample(generator) + elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": + return encoder_output.latent_dist.mode() + elif hasattr(encoder_output, "latents"): + return encoder_output.latents + else: + raise AttributeError("Could not access latents of provided encoder_output") + + +class LucyEditPipeline(DiffusionPipeline, WanLoraLoaderMixin): + r""" + Pipeline for video-to-video generation using Lucy Edit. + + This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods + implemented for all pipelines (downloading, saving, running on a particular device, etc.). + + Args: + tokenizer ([`T5Tokenizer`]): + Tokenizer from [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5Tokenizer), + specifically the [google/umt5-xxl](https://huggingface.co/google/umt5-xxl) variant. + text_encoder ([`T5EncoderModel`]): + [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically + the [google/umt5-xxl](https://huggingface.co/google/umt5-xxl) variant. + transformer ([`WanTransformer3DModel`]): + Conditional Transformer to denoise the input latents. + scheduler ([`UniPCMultistepScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded image latents. + vae ([`AutoencoderKLWan`]): + Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations. + transformer_2 ([`WanTransformer3DModel`], *optional*): + Conditional Transformer to denoise the input latents during the low-noise stage. If provided, enables + two-stage denoising where `transformer` handles high-noise stages and `transformer_2` handles low-noise + stages. If not provided, only `transformer` is used. + boundary_ratio (`float`, *optional*, defaults to `None`): + Ratio of total timesteps to use as the boundary for switching between transformers in two-stage denoising. + The actual boundary timestep is calculated as `boundary_ratio * num_train_timesteps`. When provided, + `transformer` handles timesteps >= boundary_timestep and `transformer_2` handles timesteps < + boundary_timestep. If `None`, only `transformer` is used for the entire denoising process. + """ + + model_cpu_offload_seq = "text_encoder->transformer->transformer_2->vae" + _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] + _optional_components = ["transformer", "transformer_2"] + + def __init__( + self, + tokenizer: AutoTokenizer, + text_encoder: UMT5EncoderModel, + vae: AutoencoderKLWan, + scheduler: FlowMatchEulerDiscreteScheduler, + transformer: Optional[WanTransformer3DModel] = None, + transformer_2: Optional[WanTransformer3DModel] = None, + boundary_ratio: Optional[float] = None, + expand_timesteps: bool = False, # Wan2.2 ti2v + ): + super().__init__() + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + scheduler=scheduler, + transformer_2=transformer_2, + ) + self.register_to_config(boundary_ratio=boundary_ratio) + self.register_to_config(expand_timesteps=expand_timesteps) + self.vae_scale_factor_temporal = self.vae.config.scale_factor_temporal if getattr(self, "vae", None) else 4 + self.vae_scale_factor_spatial = self.vae.config.scale_factor_spatial if getattr(self, "vae", None) else 8 + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) + + # Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline._get_t5_prompt_embeds + def _get_t5_prompt_embeds( + self, + prompt: Union[str, List[str]] = None, + num_videos_per_prompt: int = 1, + max_sequence_length: int = 226, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + prompt = [prompt_clean(u) for u in prompt] + batch_size = len(prompt) + + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask + seq_lens = mask.gt(0).sum(dim=1).long() + + prompt_embeds = self.text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] + prompt_embeds = torch.stack( + [torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0 + ) + + # duplicate text embeddings for each generation per prompt, using mps friendly method + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) + + return prompt_embeds + + # Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline.encode_prompt + def encode_prompt( + self, + prompt: Union[str, List[str]], + negative_prompt: Optional[Union[str, List[str]]] = None, + do_classifier_free_guidance: bool = True, + num_videos_per_prompt: int = 1, + prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + max_sequence_length: int = 226, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `List[str]`, *optional*): + prompt to be encoded + 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`). + do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): + Whether to use classifier free guidance or not. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + Number of videos that should be generated per prompt. torch device to place the resulting embeddings on + 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. + device: (`torch.device`, *optional*): + torch device + dtype: (`torch.dtype`, *optional*): + torch dtype + """ + device = device or self._execution_device + + prompt = [prompt] if isinstance(prompt, str) else prompt + if prompt is not None: + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + if prompt_embeds is None: + prompt_embeds = self._get_t5_prompt_embeds( + prompt=prompt, + num_videos_per_prompt=num_videos_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, + ) + + if do_classifier_free_guidance and negative_prompt_embeds is None: + negative_prompt = negative_prompt or "" + negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + + 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`." + ) + + negative_prompt_embeds = self._get_t5_prompt_embeds( + prompt=negative_prompt, + num_videos_per_prompt=num_videos_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, + ) + + return prompt_embeds, negative_prompt_embeds + + def check_inputs( + self, + video, + prompt, + negative_prompt, + height, + width, + prompt_embeds=None, + negative_prompt_embeds=None, + callback_on_step_end_tensor_inputs=None, + guidance_scale_2=None, + ): + if height % 16 != 0 or width % 16 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.") + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`: {negative_prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + elif negative_prompt is not None and ( + not isinstance(negative_prompt, str) and not isinstance(negative_prompt, list) + ): + raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}") + + if self.config.boundary_ratio is None and guidance_scale_2 is not None: + raise ValueError("`guidance_scale_2` is only supported when the pipeline's `boundary_ratio` is not None.") + + if video is None: + raise ValueError("`video` is required, received None.") + + def prepare_latents( + self, + video: Optional[torch.Tensor] = None, + batch_size: int = 1, + num_channels_latents: int = 16, + height: int = 480, + width: int = 832, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + generator: Optional[torch.Generator] = None, + latents: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + 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." + ) + + num_latent_frames = ( + (video.size(2) - 1) // self.vae_scale_factor_temporal + 1 if latents is None else latents.size(1) + ) + shape = ( + batch_size, + num_channels_latents, + num_latent_frames, + height // self.vae_scale_factor_spatial, + width // self.vae_scale_factor_spatial, + ) + # Prepare noise latents + if latents is None: + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + else: + latents = latents.to(device) + + # Prepare condition latents + condition_latents = [ + retrieve_latents(self.vae.encode(vid.unsqueeze(0)), sample_mode="argmax") for vid in video + ] + + condition_latents = torch.cat(condition_latents, dim=0).to(dtype) + + latents_mean = ( + torch.tensor(self.vae.config.latents_mean).view(1, self.vae.config.z_dim, 1, 1, 1).to(device, dtype) + ) + latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( + device, dtype + ) + + condition_latents = (condition_latents - latents_mean) * latents_std + + # Check shapes + assert latents.shape == condition_latents.shape, ( + f"Latents shape {latents.shape} does not match expected shape {condition_latents.shape}. Please check the input." + ) + + return latents, condition_latents + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1.0 + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + @property + def interrupt(self): + return self._interrupt + + @property + def attention_kwargs(self): + return self._attention_kwargs + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + video: List[Image.Image], + prompt: Union[str, List[str]] = None, + negative_prompt: Union[str, List[str]] = None, + height: int = 480, + width: int = 832, + num_frames: int = 81, + num_inference_steps: int = 50, + guidance_scale: float = 5.0, + guidance_scale_2: Optional[float] = None, + num_videos_per_prompt: Optional[int] = 1, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + output_type: Optional[str] = "np", + return_dict: bool = True, + attention_kwargs: Optional[Dict[str, Any]] = None, + callback_on_step_end: Optional[ + Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] + ] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + max_sequence_length: int = 512, + ): + r""" + The call function to the pipeline for generation. + + Args: + video (`List[Image.Image]`): + The video to use as the condition for the video generation. + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, pass `prompt_embeds` instead. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to avoid during image generation. If not defined, pass `negative_prompt_embeds` + instead. Ignored when not using guidance (`guidance_scale` < `1`). + height (`int`, defaults to `480`): + The height in pixels of the generated image. + width (`int`, defaults to `832`): + The width in pixels of the generated image. + num_frames (`int`, defaults to `81`): + The number of frames in the generated video. + num_inference_steps (`int`, defaults to `50`): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + guidance_scale (`float`, defaults to `5.0`): + Guidance scale as defined in [Classifier-Free Diffusion + Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. + of [Imagen Paper](https://huggingface.co/papers/2205.11487). 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. + guidance_scale_2 (`float`, *optional*, defaults to `None`): + Guidance scale for the low-noise stage transformer (`transformer_2`). If `None` and the pipeline's + `boundary_ratio` is not None, uses the same value as `guidance_scale`. Only used when `transformer_2` + and the pipeline's `boundary_ratio` are not None. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make + generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor is generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not + provided, text embeddings are generated from the `prompt` input argument. + output_type (`str`, *optional*, defaults to `"np"`): + The output format of the generated image. Choose between `PIL.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`LucyPipelineOutput`] instead of a plain tuple. + attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*): + A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of + each denoising step during the inference. with the following arguments: `callback_on_step_end(self: + DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a + list of all tensors as specified by `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + max_sequence_length (`int`, defaults to `512`): + The maximum sequence length of the text encoder. If the prompt is longer than this, it will be + truncated. If the prompt is shorter, it will be padded to this length. + + Examples: + + Returns: + [`~LucyPipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`LucyPipelineOutput`] is returned, otherwise a `tuple` is returned where + the first element is a list with the generated images and the second element is a list of `bool`s + indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content. + """ + + if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): + callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + video, + prompt, + negative_prompt, + height, + width, + prompt_embeds, + negative_prompt_embeds, + callback_on_step_end_tensor_inputs, + guidance_scale_2, + ) + + if num_frames % self.vae_scale_factor_temporal != 1: + logger.warning( + f"`num_frames - 1` has to be divisible by {self.vae_scale_factor_temporal}. Rounding to the nearest number." + ) + num_frames = num_frames // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1 + num_frames = max(num_frames, 1) + + if self.config.boundary_ratio is not None and guidance_scale_2 is None: + guidance_scale_2 = guidance_scale + + self._guidance_scale = guidance_scale + self._guidance_scale_2 = guidance_scale_2 + self._attention_kwargs = attention_kwargs + self._current_timestep = None + self._interrupt = False + + device = self._execution_device + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + # 3. Encode input prompt + prompt_embeds, negative_prompt_embeds = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + num_videos_per_prompt=num_videos_per_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + max_sequence_length=max_sequence_length, + device=device, + ) + + transformer_dtype = self.transformer.dtype if self.transformer is not None else self.transformer_2.dtype + prompt_embeds = prompt_embeds.to(transformer_dtype) + if negative_prompt_embeds is not None: + negative_prompt_embeds = negative_prompt_embeds.to(transformer_dtype) + + # 4. Prepare timesteps + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = self.scheduler.timesteps + + # 5. Prepare latent variables + num_channels_latents = ( + self.transformer.config.out_channels + if self.transformer is not None + else self.transformer_2.config.out_channels + ) + video = self.video_processor.preprocess_video(video, height=height, width=width).to( + device, dtype=torch.float32 + ) + latents, condition_latents = self.prepare_latents( + video, + batch_size * num_videos_per_prompt, + num_channels_latents, + height, + width, + torch.float32, + device, + generator, + latents, + ) + + mask = torch.ones(latents.shape, dtype=torch.float32, device=device) + + # 6. Denoising loop + num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + self._num_timesteps = len(timesteps) + + if self.config.boundary_ratio is not None: + boundary_timestep = self.config.boundary_ratio * self.scheduler.config.num_train_timesteps + else: + boundary_timestep = None + + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + self._current_timestep = t + + if boundary_timestep is None or t >= boundary_timestep: + # wan2.1 or high-noise stage in wan2.2 + current_model = self.transformer + current_guidance_scale = guidance_scale + else: + # low-noise stage in wan2.2 + current_model = self.transformer_2 + current_guidance_scale = guidance_scale_2 + + # latent_model_input = latents.to(transformer_dtype) + latent_model_input = torch.cat([latents, condition_latents], dim=1).to(transformer_dtype) + # latent_model_input = torch.cat([latents, latents], dim=1).to(transformer_dtype) + if self.config.expand_timesteps: + # seq_len: num_latent_frames * latent_height//2 * latent_width//2 + temp_ts = (mask[0][0][:, ::2, ::2] * t).flatten() + # batch_size, seq_len + timestep = temp_ts.unsqueeze(0).expand(latents.shape[0], -1) + else: + timestep = t.expand(latents.shape[0]) + + with current_model.cache_context("cond"): + noise_pred = current_model( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=prompt_embeds, + attention_kwargs=attention_kwargs, + return_dict=False, + )[0] + + if self.do_classifier_free_guidance: + with current_model.cache_context("uncond"): + noise_uncond = current_model( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=negative_prompt_embeds, + attention_kwargs=attention_kwargs, + return_dict=False, + )[0] + noise_pred = noise_uncond + current_guidance_scale * (noise_pred - noise_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + self._current_timestep = None + + if not output_type == "latent": + latents = latents.to(self.vae.dtype) + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( + latents.device, latents.dtype + ) + latents = latents / latents_std + latents_mean + video = self.vae.decode(latents, return_dict=False)[0] + video = self.video_processor.postprocess_video(video, output_type=output_type) + else: + video = latents + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (video,) + + return LucyPipelineOutput(frames=video) diff --git a/src/diffusers/pipelines/lucy/pipeline_output.py b/src/diffusers/pipelines/lucy/pipeline_output.py new file mode 100644 index 0000000000..cf9ea91fd1 --- /dev/null +++ b/src/diffusers/pipelines/lucy/pipeline_output.py @@ -0,0 +1,20 @@ +from dataclasses import dataclass + +import torch + +from diffusers.utils import BaseOutput + + +@dataclass +class LucyPipelineOutput(BaseOutput): + r""" + Output class for Lucy pipelines. + + Args: + frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]): + List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing + denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape + `(batch_size, num_frames, channels, height, width)`. + """ + + frames: torch.Tensor diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 00792fa55a..e29be174f0 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -1592,6 +1592,21 @@ class LTXPipeline(metaclass=DummyObject): requires_backends(cls, ["torch", "transformers"]) +class LucyEditPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class Lumina2Pipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] From d06750a5fd19781de68066bb34a3520af83cf124 Mon Sep 17 00:00:00 2001 From: Zijian Zhou Date: Wed, 17 Sep 2025 00:43:15 +0100 Subject: [PATCH 003/123] Fix autoencoder_kl_wan.py bugs for Wan2.2 VAE (#12335) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Update autoencoder_kl_wan.py When using the Wan2.2 VAE, the spatial compression ratio calculated here is incorrect. It should be 16 instead of 8. Pass it in directly via the config to ensure it’s correct here. * Update autoencoder_kl_wan.py --- src/diffusers/models/autoencoders/autoencoder_kl_wan.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_wan.py b/src/diffusers/models/autoencoders/autoencoder_kl_wan.py index d84a0861e9..e6e58c1cce 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_wan.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_wan.py @@ -1052,7 +1052,7 @@ class AutoencoderKLWan(ModelMixin, ConfigMixin, FromOriginalModelMixin): is_residual=is_residual, ) - self.spatial_compression_ratio = 2 ** len(self.temperal_downsample) + self.spatial_compression_ratio = scale_factor_spatial # When decoding a batch of video latents at a time, one can save memory by slicing across the batch dimension # to perform decoding of a single video latent at a time. @@ -1145,12 +1145,13 @@ class AutoencoderKLWan(ModelMixin, ConfigMixin, FromOriginalModelMixin): def _encode(self, x: torch.Tensor): _, _, num_frame, height, width = x.shape - if self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height): - return self.tiled_encode(x) - self.clear_cache() if self.config.patch_size is not None: x = patchify(x, patch_size=self.config.patch_size) + + if self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height): + return self.tiled_encode(x) + iter_ = 1 + (num_frame - 1) // 4 for i in range(iter_): self._enc_conv_idx = [0] From efb7a299af46d739dec6a57a5d2814165fba24b5 Mon Sep 17 00:00:00 2001 From: DefTruth <31974251+DefTruth@users.noreply.github.com> Date: Wed, 17 Sep 2025 12:52:15 +0800 Subject: [PATCH 004/123] Fix many type hint errors (#12289) * fix hidream type hint * fix hunyuan-video type hint * fix many type hint * fix many type hint errors * fix many type hint errors * fix many type hint errors * make stype & make quality --- src/diffusers/models/attention.py | 2 +- .../models/transformers/auraflow_transformer_2d.py | 10 +++++----- .../models/transformers/cogvideox_transformer_3d.py | 4 ++-- .../models/transformers/consisid_transformer_3d.py | 4 ++-- src/diffusers/models/transformers/lumina_nextdit2d.py | 6 +++--- src/diffusers/models/transformers/transformer_bria.py | 4 ++-- .../models/transformers/transformer_cogview3plus.py | 6 +++--- .../models/transformers/transformer_cogview4.py | 4 ++-- .../models/transformers/transformer_hidream_image.py | 10 +++++----- .../models/transformers/transformer_hunyuan_video.py | 6 +++--- .../transformer_hunyuan_video_framepack.py | 4 ++-- 11 files changed, 30 insertions(+), 30 deletions(-) diff --git a/src/diffusers/models/attention.py b/src/diffusers/models/attention.py index c720b37955..c99133f257 100644 --- a/src/diffusers/models/attention.py +++ b/src/diffusers/models/attention.py @@ -674,7 +674,7 @@ class JointTransformerBlock(nn.Module): encoder_hidden_states: torch.FloatTensor, temb: torch.FloatTensor, joint_attention_kwargs: Optional[Dict[str, Any]] = None, - ): + ) -> Tuple[torch.Tensor, torch.Tensor]: joint_attention_kwargs = joint_attention_kwargs or {} if self.use_dual_attention: norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp, norm_hidden_states2, gate_msa2 = self.norm1( diff --git a/src/diffusers/models/transformers/auraflow_transformer_2d.py b/src/diffusers/models/transformers/auraflow_transformer_2d.py index a8d275d142..4d7d1ba40e 100644 --- a/src/diffusers/models/transformers/auraflow_transformer_2d.py +++ b/src/diffusers/models/transformers/auraflow_transformer_2d.py @@ -13,7 +13,7 @@ # limitations under the License. -from typing import Any, Dict, Optional, Union +from typing import Any, Dict, Optional, Tuple, Union import torch import torch.nn as nn @@ -92,7 +92,7 @@ class AuraFlowPatchEmbed(nn.Module): return selected_indices - def forward(self, latent): + def forward(self, latent) -> torch.Tensor: batch_size, num_channels, height, width = latent.size() latent = latent.view( batch_size, @@ -173,7 +173,7 @@ class AuraFlowSingleTransformerBlock(nn.Module): hidden_states: torch.FloatTensor, temb: torch.FloatTensor, attention_kwargs: Optional[Dict[str, Any]] = None, - ): + ) -> torch.Tensor: residual = hidden_states attention_kwargs = attention_kwargs or {} @@ -242,7 +242,7 @@ class AuraFlowJointTransformerBlock(nn.Module): encoder_hidden_states: torch.FloatTensor, temb: torch.FloatTensor, attention_kwargs: Optional[Dict[str, Any]] = None, - ): + ) -> Tuple[torch.Tensor, torch.Tensor]: residual = hidden_states residual_context = encoder_hidden_states attention_kwargs = attention_kwargs or {} @@ -472,7 +472,7 @@ class AuraFlowTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, From timestep: torch.LongTensor = None, attention_kwargs: Optional[Dict[str, Any]] = None, return_dict: bool = True, - ) -> Union[torch.FloatTensor, Transformer2DModelOutput]: + ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]: if attention_kwargs is not None: attention_kwargs = attention_kwargs.copy() lora_scale = attention_kwargs.pop("scale", 1.0) diff --git a/src/diffusers/models/transformers/cogvideox_transformer_3d.py b/src/diffusers/models/transformers/cogvideox_transformer_3d.py index a8c98bccb8..5038109690 100644 --- a/src/diffusers/models/transformers/cogvideox_transformer_3d.py +++ b/src/diffusers/models/transformers/cogvideox_transformer_3d.py @@ -122,7 +122,7 @@ class CogVideoXBlock(nn.Module): temb: torch.Tensor, image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> torch.Tensor: + ) -> Tuple[torch.Tensor, torch.Tensor]: text_seq_length = encoder_hidden_states.size(1) attention_kwargs = attention_kwargs or {} @@ -441,7 +441,7 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, Cac image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, attention_kwargs: Optional[Dict[str, Any]] = None, return_dict: bool = True, - ): + ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]: if attention_kwargs is not None: attention_kwargs = attention_kwargs.copy() lora_scale = attention_kwargs.pop("scale", 1.0) diff --git a/src/diffusers/models/transformers/consisid_transformer_3d.py b/src/diffusers/models/transformers/consisid_transformer_3d.py index 41632dbd47..91fe811f00 100644 --- a/src/diffusers/models/transformers/consisid_transformer_3d.py +++ b/src/diffusers/models/transformers/consisid_transformer_3d.py @@ -315,7 +315,7 @@ class ConsisIDBlock(nn.Module): encoder_hidden_states: torch.Tensor, temb: torch.Tensor, image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - ) -> torch.Tensor: + ) -> Tuple[torch.Tensor, torch.Tensor]: text_seq_length = encoder_hidden_states.size(1) # norm & modulate @@ -691,7 +691,7 @@ class ConsisIDTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin): id_cond: Optional[torch.Tensor] = None, id_vit_hidden: Optional[torch.Tensor] = None, return_dict: bool = True, - ): + ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]: if attention_kwargs is not None: attention_kwargs = attention_kwargs.copy() lora_scale = attention_kwargs.pop("scale", 1.0) diff --git a/src/diffusers/models/transformers/lumina_nextdit2d.py b/src/diffusers/models/transformers/lumina_nextdit2d.py index 84b1175386..bed5e69c2d 100644 --- a/src/diffusers/models/transformers/lumina_nextdit2d.py +++ b/src/diffusers/models/transformers/lumina_nextdit2d.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Dict, Optional +from typing import Any, Dict, Optional, Tuple, Union import torch import torch.nn as nn @@ -124,7 +124,7 @@ class LuminaNextDiTBlock(nn.Module): encoder_mask: torch.Tensor, temb: torch.Tensor, cross_attention_kwargs: Optional[Dict[str, Any]] = None, - ): + ) -> torch.Tensor: """ Perform a forward pass through the LuminaNextDiTBlock. @@ -297,7 +297,7 @@ class LuminaNextDiT2DModel(ModelMixin, ConfigMixin): image_rotary_emb: torch.Tensor, cross_attention_kwargs: Dict[str, Any] = None, return_dict=True, - ) -> torch.Tensor: + ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]: """ Forward pass of LuminaNextDiT. diff --git a/src/diffusers/models/transformers/transformer_bria.py b/src/diffusers/models/transformers/transformer_bria.py index 27a9941501..04a9c5645c 100644 --- a/src/diffusers/models/transformers/transformer_bria.py +++ b/src/diffusers/models/transformers/transformer_bria.py @@ -472,7 +472,7 @@ class BriaSingleTransformerBlock(nn.Module): temb: torch.Tensor, image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> torch.Tensor: + ) -> Tuple[torch.Tensor, torch.Tensor]: text_seq_len = encoder_hidden_states.shape[1] hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) @@ -588,7 +588,7 @@ class BriaTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOrig return_dict: bool = True, controlnet_block_samples=None, controlnet_single_block_samples=None, - ) -> Union[torch.FloatTensor, Transformer2DModelOutput]: + ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]: """ The [`BriaTransformer2DModel`] forward method. diff --git a/src/diffusers/models/transformers/transformer_cogview3plus.py b/src/diffusers/models/transformers/transformer_cogview3plus.py index 77f15f6ca6..7356f4a606 100644 --- a/src/diffusers/models/transformers/transformer_cogview3plus.py +++ b/src/diffusers/models/transformers/transformer_cogview3plus.py @@ -13,7 +13,7 @@ # limitations under the License. -from typing import Dict, Union +from typing import Dict, Tuple, Union import torch import torch.nn as nn @@ -79,7 +79,7 @@ class CogView3PlusTransformerBlock(nn.Module): hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, emb: torch.Tensor, - ) -> torch.Tensor: + ) -> Tuple[torch.Tensor, torch.Tensor]: text_seq_length = encoder_hidden_states.size(1) # norm & modulate @@ -293,7 +293,7 @@ class CogView3PlusTransformer2DModel(ModelMixin, ConfigMixin): target_size: torch.Tensor, crop_coords: torch.Tensor, return_dict: bool = True, - ) -> Union[torch.Tensor, Transformer2DModelOutput]: + ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]: """ The [`CogView3PlusTransformer2DModel`] forward method. diff --git a/src/diffusers/models/transformers/transformer_cogview4.py b/src/diffusers/models/transformers/transformer_cogview4.py index 25dcfa14cc..64e9a538a7 100644 --- a/src/diffusers/models/transformers/transformer_cogview4.py +++ b/src/diffusers/models/transformers/transformer_cogview4.py @@ -494,7 +494,7 @@ class CogView4TransformerBlock(nn.Module): ] = None, attention_mask: Optional[Dict[str, torch.Tensor]] = None, attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> torch.Tensor: + ) -> Tuple[torch.Tensor, torch.Tensor]: # 1. Timestep conditioning ( norm_hidden_states, @@ -717,7 +717,7 @@ class CogView4Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, Cach image_rotary_emb: Optional[ Union[Tuple[torch.Tensor, torch.Tensor], List[Tuple[torch.Tensor, torch.Tensor]]] ] = None, - ) -> Union[torch.Tensor, Transformer2DModelOutput]: + ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]: if attention_kwargs is not None: attention_kwargs = attention_kwargs.copy() lora_scale = attention_kwargs.pop("scale", 1.0) diff --git a/src/diffusers/models/transformers/transformer_hidream_image.py b/src/diffusers/models/transformers/transformer_hidream_image.py index 77902dcf58..4a5aee29ab 100644 --- a/src/diffusers/models/transformers/transformer_hidream_image.py +++ b/src/diffusers/models/transformers/transformer_hidream_image.py @@ -55,7 +55,7 @@ class HiDreamImageTimestepEmbed(nn.Module): self.time_proj = Timesteps(num_channels=frequency_embedding_size, flip_sin_to_cos=True, downscale_freq_shift=0) self.timestep_embedder = TimestepEmbedding(in_channels=frequency_embedding_size, time_embed_dim=hidden_size) - def forward(self, timesteps: torch.Tensor, wdtype: Optional[torch.dtype] = None): + def forward(self, timesteps: torch.Tensor, wdtype: Optional[torch.dtype] = None) -> torch.Tensor: t_emb = self.time_proj(timesteps).to(dtype=wdtype) t_emb = self.timestep_embedder(t_emb) return t_emb @@ -87,7 +87,7 @@ class HiDreamImagePatchEmbed(nn.Module): self.out_channels = out_channels self.proj = nn.Linear(in_channels * patch_size * patch_size, out_channels, bias=True) - def forward(self, latent): + def forward(self, latent) -> torch.Tensor: latent = self.proj(latent) return latent @@ -534,7 +534,7 @@ class HiDreamImageTransformerBlock(nn.Module): encoder_hidden_states: Optional[torch.Tensor] = None, temb: Optional[torch.Tensor] = None, image_rotary_emb: torch.Tensor = None, - ) -> torch.Tensor: + ) -> Tuple[torch.Tensor, torch.Tensor]: wtype = hidden_states.dtype ( shift_msa_i, @@ -592,7 +592,7 @@ class HiDreamBlock(nn.Module): encoder_hidden_states: Optional[torch.Tensor] = None, temb: Optional[torch.Tensor] = None, image_rotary_emb: torch.Tensor = None, - ) -> torch.Tensor: + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: return self.block( hidden_states=hidden_states, hidden_states_masks=hidden_states_masks, @@ -786,7 +786,7 @@ class HiDreamImageTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, attention_kwargs: Optional[Dict[str, Any]] = None, return_dict: bool = True, **kwargs, - ): + ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]: encoder_hidden_states = kwargs.get("encoder_hidden_states", None) if encoder_hidden_states is not None: diff --git a/src/diffusers/models/transformers/transformer_hunyuan_video.py b/src/diffusers/models/transformers/transformer_hunyuan_video.py index 6944a6c536..bc857ccab4 100644 --- a/src/diffusers/models/transformers/transformer_hunyuan_video.py +++ b/src/diffusers/models/transformers/transformer_hunyuan_video.py @@ -529,7 +529,7 @@ class HunyuanVideoSingleTransformerBlock(nn.Module): image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, *args, **kwargs, - ) -> torch.Tensor: + ) -> Tuple[torch.Tensor, torch.Tensor]: text_seq_length = encoder_hidden_states.shape[1] hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1) @@ -684,7 +684,7 @@ class HunyuanVideoTokenReplaceSingleTransformerBlock(nn.Module): image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, token_replace_emb: torch.Tensor = None, num_tokens: int = None, - ) -> torch.Tensor: + ) -> Tuple[torch.Tensor, torch.Tensor]: text_seq_length = encoder_hidden_states.shape[1] hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1) @@ -1038,7 +1038,7 @@ class HunyuanVideoTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, guidance: torch.Tensor = None, attention_kwargs: Optional[Dict[str, Any]] = None, return_dict: bool = True, - ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: + ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]: if attention_kwargs is not None: attention_kwargs = attention_kwargs.copy() lora_scale = attention_kwargs.pop("scale", 1.0) diff --git a/src/diffusers/models/transformers/transformer_hunyuan_video_framepack.py b/src/diffusers/models/transformers/transformer_hunyuan_video_framepack.py index c2eb7fd2a7..60b40fff3c 100644 --- a/src/diffusers/models/transformers/transformer_hunyuan_video_framepack.py +++ b/src/diffusers/models/transformers/transformer_hunyuan_video_framepack.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple, Union import torch import torch.nn as nn @@ -216,7 +216,7 @@ class HunyuanVideoFramepackTransformer3DModel( indices_latents_history_4x: Optional[torch.Tensor] = None, attention_kwargs: Optional[Dict[str, Any]] = None, return_dict: bool = True, - ): + ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]: if attention_kwargs is not None: attention_kwargs = attention_kwargs.copy() lora_scale = attention_kwargs.pop("scale", 1.0) From eda9ff8300eb3b8ceec15ef69d74e35abd3d39b3 Mon Sep 17 00:00:00 2001 From: Fredy Date: Thu, 18 Sep 2025 00:03:43 -0600 Subject: [PATCH 005/123] Add RequestScopedPipeline for safe concurrent inference, tokenizer lock and non-mutating retrieve_timesteps (#12328) * Basic implementation of request scheduling * Basic editing in SD and Flux Pipelines * Small Fix * Fix * Update for more pipelines * Add examples/server-async * Add examples/server-async * Updated RequestScopedPipeline to handle a single tokenizer lock to avoid race conditions * Fix * Fix _TokenizerLockWrapper * Fix _TokenizerLockWrapper * Delete _TokenizerLockWrapper * Fix tokenizer * Update examples/server-async * Fix server-async * Optimizations in examples/server-async * We keep the implementation simple in examples/server-async * Update examples/server-async/README.md * Update examples/server-async/README.md for changes to tokenizer locks and backward-compatible retrieve_timesteps * The changes to the diffusers core have been undone and all logic is being moved to exmaples/server-async * Update examples/server-async/utils/* * Fix BaseAsyncScheduler * Rollback in the core of the diffusers * Update examples/server-async/README.md * Complete rollback of diffusers core files * Simple implementation of an asynchronous server compatible with SD3-3.5 and Flux Pipelines * Update examples/server-async/README.md * Fixed import errors in 'examples/server-async/serverasync.py' * Flux Pipeline Discard * Update examples/server-async/README.md * Apply style fixes --------- Co-authored-by: Sayak Paul Co-authored-by: github-actions[bot] --- examples/server-async/Pipelines.py | 91 ++++++ examples/server-async/README.md | 171 ++++++++++ examples/server-async/requirements.txt | 10 + examples/server-async/serverasync.py | 230 ++++++++++++++ examples/server-async/test.py | 65 ++++ examples/server-async/utils/__init__.py | 2 + .../utils/requestscopedpipeline.py | 296 ++++++++++++++++++ examples/server-async/utils/scheduler.py | 141 +++++++++ examples/server-async/utils/utils.py | 48 +++ 9 files changed, 1054 insertions(+) create mode 100644 examples/server-async/Pipelines.py create mode 100644 examples/server-async/README.md create mode 100644 examples/server-async/requirements.txt create mode 100644 examples/server-async/serverasync.py create mode 100644 examples/server-async/test.py create mode 100644 examples/server-async/utils/__init__.py create mode 100644 examples/server-async/utils/requestscopedpipeline.py create mode 100644 examples/server-async/utils/scheduler.py create mode 100644 examples/server-async/utils/utils.py diff --git a/examples/server-async/Pipelines.py b/examples/server-async/Pipelines.py new file mode 100644 index 0000000000..f89cac6a7e --- /dev/null +++ b/examples/server-async/Pipelines.py @@ -0,0 +1,91 @@ +import logging +import os +from dataclasses import dataclass, field +from typing import List + +import torch +from pydantic import BaseModel + +from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3 import StableDiffusion3Pipeline + + +logger = logging.getLogger(__name__) + + +class TextToImageInput(BaseModel): + model: str + prompt: str + size: str | None = None + n: int | None = None + + +@dataclass +class PresetModels: + SD3: List[str] = field(default_factory=lambda: ["stabilityai/stable-diffusion-3-medium"]) + SD3_5: List[str] = field( + default_factory=lambda: [ + "stabilityai/stable-diffusion-3.5-large", + "stabilityai/stable-diffusion-3.5-large-turbo", + "stabilityai/stable-diffusion-3.5-medium", + ] + ) + + +class TextToImagePipelineSD3: + def __init__(self, model_path: str | None = None): + self.model_path = model_path or os.getenv("MODEL_PATH") + self.pipeline: StableDiffusion3Pipeline | None = None + self.device: str | None = None + + def start(self): + if torch.cuda.is_available(): + model_path = self.model_path or "stabilityai/stable-diffusion-3.5-large" + logger.info("Loading CUDA") + self.device = "cuda" + self.pipeline = StableDiffusion3Pipeline.from_pretrained( + model_path, + torch_dtype=torch.float16, + ).to(device=self.device) + elif torch.backends.mps.is_available(): + model_path = self.model_path or "stabilityai/stable-diffusion-3.5-medium" + logger.info("Loading MPS for Mac M Series") + self.device = "mps" + self.pipeline = StableDiffusion3Pipeline.from_pretrained( + model_path, + torch_dtype=torch.bfloat16, + ).to(device=self.device) + else: + raise Exception("No CUDA or MPS device available") + + +class ModelPipelineInitializer: + def __init__(self, model: str = "", type_models: str = "t2im"): + self.model = model + self.type_models = type_models + self.pipeline = None + self.device = "cuda" if torch.cuda.is_available() else "mps" + self.model_type = None + + def initialize_pipeline(self): + if not self.model: + raise ValueError("Model name not provided") + + # Check if model exists in PresetModels + preset_models = PresetModels() + + # Determine which model type we're dealing with + if self.model in preset_models.SD3: + self.model_type = "SD3" + elif self.model in preset_models.SD3_5: + self.model_type = "SD3_5" + + # Create appropriate pipeline based on model type and type_models + if self.type_models == "t2im": + if self.model_type in ["SD3", "SD3_5"]: + self.pipeline = TextToImagePipelineSD3(self.model) + else: + raise ValueError(f"Model type {self.model_type} not supported for text-to-image") + elif self.type_models == "t2v": + raise ValueError(f"Unsupported type_models: {self.type_models}") + + return self.pipeline diff --git a/examples/server-async/README.md b/examples/server-async/README.md new file mode 100644 index 0000000000..a47ab7c7f2 --- /dev/null +++ b/examples/server-async/README.md @@ -0,0 +1,171 @@ +# Asynchronous server and parallel execution of models + +> Example/demo server that keeps a single model in memory while safely running parallel inference requests by creating per-request lightweight views and cloning only small, stateful components (schedulers, RNG state, small mutable attrs). Works with StableDiffusion3 pipelines. +> We recommend running 10 to 50 inferences in parallel for optimal performance, averaging between 25 and 30 seconds to 1 minute and 1 minute and 30 seconds. (This is only recommended if you have a GPU with 35GB of VRAM or more; otherwise, keep it to one or two inferences in parallel to avoid decoding or saving errors due to memory shortages.) + +## ⚠️ IMPORTANT + +* The example demonstrates how to run pipelines like `StableDiffusion3-3.5` concurrently while keeping a single copy of the heavy model parameters on GPU. + +## Necessary components + +All the components needed to create the inference server are in the current directory: + +``` +server-async/ +├── utils/ +├─────── __init__.py +├─────── scheduler.py # BaseAsyncScheduler wrapper and async_retrieve_timesteps for secure inferences +├─────── requestscopedpipeline.py # RequestScoped Pipeline for inference with a single in-memory model +├─────── utils.py # Image/video saving utilities and service configuration +├── Pipelines.py # pipeline loader classes (SD3) +├── serverasync.py # FastAPI app with lifespan management and async inference endpoints +├── test.py # Client test script for inference requests +├── requirements.txt # Dependencies +└── README.md # This documentation +``` + +## What `diffusers-async` adds / Why we needed it + +Core problem: a naive server that calls `pipe.__call__` concurrently can hit **race conditions** (e.g., `scheduler.set_timesteps` mutates shared state) or explode memory by deep-copying the whole pipeline per-request. + +`diffusers-async` / this example addresses that by: + +* **Request-scoped views**: `RequestScopedPipeline` creates a shallow copy of the pipeline per request so heavy weights (UNet, VAE, text encoder) remain shared and *are not duplicated*. +* **Per-request mutable state**: stateful small objects (scheduler, RNG state, small lists/dicts, callbacks) are cloned per request. The system uses `BaseAsyncScheduler.clone_for_request(...)` for scheduler cloning, with fallback to safe `deepcopy` or other heuristics. +* **Tokenizer concurrency safety**: `RequestScopedPipeline` now manages an internal tokenizer lock with automatic tokenizer detection and wrapping. This ensures that Rust tokenizers are safe to use under concurrency — race condition errors like `Already borrowed` no longer occur. +* **`async_retrieve_timesteps(..., return_scheduler=True)`**: fully retro-compatible helper that returns `(timesteps, num_inference_steps, scheduler)` without mutating the shared scheduler. For users not using `return_scheduler=True`, the behavior is identical to the original API. +* **Robust attribute handling**: wrapper avoids writing to read-only properties (e.g., `components`) and auto-detects small mutable attributes to clone while avoiding duplication of large tensors. Configurable tensor size threshold prevents cloning of large tensors. +* **Enhanced scheduler wrapping**: `BaseAsyncScheduler` automatically wraps schedulers with improved `__getattr__`, `__setattr__`, and debugging methods (`__repr__`, `__str__`). + +## How the server works (high-level flow) + +1. **Single model instance** is loaded into memory (GPU/MPS) when the server starts. +2. On each HTTP inference request: + + * The server uses `RequestScopedPipeline.generate(...)` which: + + * automatically wraps the base scheduler in `BaseAsyncScheduler` (if not already wrapped), + * obtains a *local scheduler* (via `clone_for_request()` or `deepcopy`), + * does `local_pipe = copy.copy(base_pipe)` (shallow copy), + * sets `local_pipe.scheduler = local_scheduler` (if possible), + * clones only small mutable attributes (callbacks, rng, small latents) with auto-detection, + * wraps tokenizers with thread-safe locks to prevent race conditions, + * optionally enters a `model_cpu_offload_context()` for memory offload hooks, + * calls the pipeline on the local view (`local_pipe(...)`). +3. **Result**: inference completes, images are moved to CPU & saved (if requested), internal buffers freed (GC + `torch.cuda.empty_cache()`). +4. Multiple requests can run in parallel while sharing heavy weights and isolating mutable state. + +## How to set up and run the server + +### 1) Install dependencies + +Recommended: create a virtualenv / conda environment. + +```bash +pip install diffusers +pip install -r requirements.txt +``` + +### 2) Start the server + +Using the `serverasync.py` file that already has everything you need: + +```bash +python serverasync.py +``` + +The server will start on `http://localhost:8500` by default with the following features: +- FastAPI application with async lifespan management +- Automatic model loading and pipeline initialization +- Request counting and active inference tracking +- Memory cleanup after each inference +- CORS middleware for cross-origin requests + +### 3) Test the server + +Use the included test script: + +```bash +python test.py +``` + +Or send a manual request: + +`POST /api/diffusers/inference` with JSON body: + +```json +{ + "prompt": "A futuristic cityscape, vibrant colors", + "num_inference_steps": 30, + "num_images_per_prompt": 1 +} +``` + +Response example: + +```json +{ + "response": ["http://localhost:8500/images/img123.png"] +} +``` + +### 4) Server endpoints + +- `GET /` - Welcome message +- `POST /api/diffusers/inference` - Main inference endpoint +- `GET /images/{filename}` - Serve generated images +- `GET /api/status` - Server status and memory info + +## Advanced Configuration + +### RequestScopedPipeline Parameters + +```python +RequestScopedPipeline( + pipeline, # Base pipeline to wrap + mutable_attrs=None, # Custom list of attributes to clone + auto_detect_mutables=True, # Enable automatic detection of mutable attributes + tensor_numel_threshold=1_000_000, # Tensor size threshold for cloning + tokenizer_lock=None, # Custom threading lock for tokenizers + wrap_scheduler=True # Auto-wrap scheduler in BaseAsyncScheduler +) +``` + +### BaseAsyncScheduler Features + +* Transparent proxy to the original scheduler with `__getattr__` and `__setattr__` +* `clone_for_request()` method for safe per-request scheduler cloning +* Enhanced debugging with `__repr__` and `__str__` methods +* Full compatibility with existing scheduler APIs + +### Server Configuration + +The server configuration can be modified in `serverasync.py` through the `ServerConfigModels` dataclass: + +```python +@dataclass +class ServerConfigModels: + model: str = 'stabilityai/stable-diffusion-3.5-medium' + type_models: str = 't2im' + host: str = '0.0.0.0' + port: int = 8500 +``` + +## Troubleshooting (quick) + +* `Already borrowed` — previously a Rust tokenizer concurrency error. + ✅ This is now fixed: `RequestScopedPipeline` automatically detects and wraps tokenizers with thread locks, so race conditions no longer happen. + +* `can't set attribute 'components'` — pipeline exposes read-only `components`. + ✅ The RequestScopedPipeline now detects read-only properties and skips setting them automatically. + +* Scheduler issues: + * If the scheduler doesn't implement `clone_for_request` and `deepcopy` fails, we log and fallback — but prefer `async_retrieve_timesteps(..., return_scheduler=True)` to avoid mutating the shared scheduler. + ✅ Note: `async_retrieve_timesteps` is fully retro-compatible — if you don't pass `return_scheduler=True`, the behavior is unchanged. + +* Memory issues with large tensors: + ✅ The system now has configurable `tensor_numel_threshold` to prevent cloning of large tensors while still cloning small mutable ones. + +* Automatic tokenizer detection: + ✅ The system automatically identifies tokenizer components by checking for tokenizer methods, class names, and attributes, then applies thread-safe wrappers. \ No newline at end of file diff --git a/examples/server-async/requirements.txt b/examples/server-async/requirements.txt new file mode 100644 index 0000000000..aafa93b702 --- /dev/null +++ b/examples/server-async/requirements.txt @@ -0,0 +1,10 @@ +torch +torchvision +transformers +sentencepiece +fastapi +uvicorn +ftfy +accelerate +xformers +protobuf \ No newline at end of file diff --git a/examples/server-async/serverasync.py b/examples/server-async/serverasync.py new file mode 100644 index 0000000000..b279b36f9a --- /dev/null +++ b/examples/server-async/serverasync.py @@ -0,0 +1,230 @@ +import asyncio +import gc +import logging +import os +import random +import threading +from contextlib import asynccontextmanager +from dataclasses import dataclass +from typing import Any, Dict, Optional, Type + +import torch +from fastapi import FastAPI, HTTPException, Request +from fastapi.concurrency import run_in_threadpool +from fastapi.middleware.cors import CORSMiddleware +from fastapi.responses import FileResponse +from Pipelines import ModelPipelineInitializer +from pydantic import BaseModel + +from utils import RequestScopedPipeline, Utils + + +@dataclass +class ServerConfigModels: + model: str = "stabilityai/stable-diffusion-3.5-medium" + type_models: str = "t2im" + constructor_pipeline: Optional[Type] = None + custom_pipeline: Optional[Type] = None + components: Optional[Dict[str, Any]] = None + torch_dtype: Optional[torch.dtype] = None + host: str = "0.0.0.0" + port: int = 8500 + + +server_config = ServerConfigModels() + + +@asynccontextmanager +async def lifespan(app: FastAPI): + logging.basicConfig(level=logging.INFO) + app.state.logger = logging.getLogger("diffusers-server") + os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128,expandable_segments:True" + os.environ["CUDA_LAUNCH_BLOCKING"] = "0" + + app.state.total_requests = 0 + app.state.active_inferences = 0 + app.state.metrics_lock = asyncio.Lock() + app.state.metrics_task = None + + app.state.utils_app = Utils( + host=server_config.host, + port=server_config.port, + ) + + async def metrics_loop(): + try: + while True: + async with app.state.metrics_lock: + total = app.state.total_requests + active = app.state.active_inferences + app.state.logger.info(f"[METRICS] total_requests={total} active_inferences={active}") + await asyncio.sleep(5) + except asyncio.CancelledError: + app.state.logger.info("Metrics loop cancelled") + raise + + app.state.metrics_task = asyncio.create_task(metrics_loop()) + + try: + yield + finally: + task = app.state.metrics_task + if task: + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + try: + stop_fn = getattr(model_pipeline, "stop", None) or getattr(model_pipeline, "close", None) + if callable(stop_fn): + await run_in_threadpool(stop_fn) + except Exception as e: + app.state.logger.warning(f"Error during pipeline shutdown: {e}") + + app.state.logger.info("Lifespan shutdown complete") + + +app = FastAPI(lifespan=lifespan) + +logger = logging.getLogger("DiffusersServer.Pipelines") + + +initializer = ModelPipelineInitializer( + model=server_config.model, + type_models=server_config.type_models, +) +model_pipeline = initializer.initialize_pipeline() +model_pipeline.start() + +request_pipe = RequestScopedPipeline(model_pipeline.pipeline) +pipeline_lock = threading.Lock() + +logger.info(f"Pipeline initialized and ready to receive requests (model ={server_config.model})") + +app.state.MODEL_INITIALIZER = initializer +app.state.MODEL_PIPELINE = model_pipeline +app.state.REQUEST_PIPE = request_pipe +app.state.PIPELINE_LOCK = pipeline_lock + + +class JSONBodyQueryAPI(BaseModel): + model: str | None = None + prompt: str + negative_prompt: str | None = None + num_inference_steps: int = 28 + num_images_per_prompt: int = 1 + + +@app.middleware("http") +async def count_requests_middleware(request: Request, call_next): + async with app.state.metrics_lock: + app.state.total_requests += 1 + response = await call_next(request) + return response + + +@app.get("/") +async def root(): + return {"message": "Welcome to the Diffusers Server"} + + +@app.post("/api/diffusers/inference") +async def api(json: JSONBodyQueryAPI): + prompt = json.prompt + negative_prompt = json.negative_prompt or "" + num_steps = json.num_inference_steps + num_images_per_prompt = json.num_images_per_prompt + + wrapper = app.state.MODEL_PIPELINE + initializer = app.state.MODEL_INITIALIZER + + utils_app = app.state.utils_app + + if not wrapper or not wrapper.pipeline: + raise HTTPException(500, "Model not initialized correctly") + if not prompt.strip(): + raise HTTPException(400, "No prompt provided") + + def make_generator(): + g = torch.Generator(device=initializer.device) + return g.manual_seed(random.randint(0, 10_000_000)) + + req_pipe = app.state.REQUEST_PIPE + + def infer(): + gen = make_generator() + return req_pipe.generate( + prompt=prompt, + negative_prompt=negative_prompt, + generator=gen, + num_inference_steps=num_steps, + num_images_per_prompt=num_images_per_prompt, + device=initializer.device, + output_type="pil", + ) + + try: + async with app.state.metrics_lock: + app.state.active_inferences += 1 + + output = await run_in_threadpool(infer) + + async with app.state.metrics_lock: + app.state.active_inferences = max(0, app.state.active_inferences - 1) + + urls = [utils_app.save_image(img) for img in output.images] + return {"response": urls} + + except Exception as e: + async with app.state.metrics_lock: + app.state.active_inferences = max(0, app.state.active_inferences - 1) + logger.error(f"Error during inference: {e}") + raise HTTPException(500, f"Error in processing: {e}") + + finally: + if torch.cuda.is_available(): + torch.cuda.synchronize() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + torch.cuda.ipc_collect() + gc.collect() + + +@app.get("/images/{filename}") +async def serve_image(filename: str): + utils_app = app.state.utils_app + file_path = os.path.join(utils_app.image_dir, filename) + if not os.path.isfile(file_path): + raise HTTPException(status_code=404, detail="Image not found") + return FileResponse(file_path, media_type="image/png") + + +@app.get("/api/status") +async def get_status(): + memory_info = {} + if torch.cuda.is_available(): + memory_allocated = torch.cuda.memory_allocated() / 1024**3 # GB + memory_reserved = torch.cuda.memory_reserved() / 1024**3 # GB + memory_info = { + "memory_allocated_gb": round(memory_allocated, 2), + "memory_reserved_gb": round(memory_reserved, 2), + "device": torch.cuda.get_device_name(0), + } + + return {"current_model": server_config.model, "type_models": server_config.type_models, "memory": memory_info} + + +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + +if __name__ == "__main__": + import uvicorn + + uvicorn.run(app, host=server_config.host, port=server_config.port) diff --git a/examples/server-async/test.py b/examples/server-async/test.py new file mode 100644 index 0000000000..e67317ea8f --- /dev/null +++ b/examples/server-async/test.py @@ -0,0 +1,65 @@ +import os +import time +import urllib.parse + +import requests + + +SERVER_URL = "http://localhost:8500/api/diffusers/inference" +BASE_URL = "http://localhost:8500" +DOWNLOAD_FOLDER = "generated_images" +WAIT_BEFORE_DOWNLOAD = 2 # seconds + +os.makedirs(DOWNLOAD_FOLDER, exist_ok=True) + + +def save_from_url(url: str) -> str: + """Download the given URL (relative or absolute) and save it locally.""" + if url.startswith("/"): + direct = BASE_URL.rstrip("/") + url + else: + direct = url + resp = requests.get(direct, timeout=60) + resp.raise_for_status() + filename = os.path.basename(urllib.parse.urlparse(direct).path) or f"img_{int(time.time())}.png" + path = os.path.join(DOWNLOAD_FOLDER, filename) + with open(path, "wb") as f: + f.write(resp.content) + return path + + +def main(): + payload = { + "prompt": "The T-800 Terminator Robot Returning From The Future, Anime Style", + "num_inference_steps": 30, + "num_images_per_prompt": 1, + } + + print("Sending request...") + try: + r = requests.post(SERVER_URL, json=payload, timeout=480) + r.raise_for_status() + except Exception as e: + print(f"Request failed: {e}") + return + + body = r.json().get("response", []) + # Normalize to a list + urls = body if isinstance(body, list) else [body] if body else [] + if not urls: + print("No URLs found in the response. Check the server output.") + return + + print(f"Received {len(urls)} URL(s). Waiting {WAIT_BEFORE_DOWNLOAD}s before downloading...") + time.sleep(WAIT_BEFORE_DOWNLOAD) + + for u in urls: + try: + path = save_from_url(u) + print(f"Image saved to: {path}") + except Exception as e: + print(f"Error downloading {u}: {e}") + + +if __name__ == "__main__": + main() diff --git a/examples/server-async/utils/__init__.py b/examples/server-async/utils/__init__.py new file mode 100644 index 0000000000..731cfe491a --- /dev/null +++ b/examples/server-async/utils/__init__.py @@ -0,0 +1,2 @@ +from .requestscopedpipeline import RequestScopedPipeline +from .utils import Utils diff --git a/examples/server-async/utils/requestscopedpipeline.py b/examples/server-async/utils/requestscopedpipeline.py new file mode 100644 index 0000000000..57d1e25671 --- /dev/null +++ b/examples/server-async/utils/requestscopedpipeline.py @@ -0,0 +1,296 @@ +import copy +import threading +from typing import Any, Iterable, List, Optional + +import torch + +from diffusers.utils import logging + +from .scheduler import BaseAsyncScheduler, async_retrieve_timesteps + + +logger = logging.get_logger(__name__) + + +def safe_tokenize(tokenizer, *args, lock, **kwargs): + with lock: + return tokenizer(*args, **kwargs) + + +class RequestScopedPipeline: + DEFAULT_MUTABLE_ATTRS = [ + "_all_hooks", + "_offload_device", + "_progress_bar_config", + "_progress_bar", + "_rng_state", + "_last_seed", + "latents", + ] + + def __init__( + self, + pipeline: Any, + mutable_attrs: Optional[Iterable[str]] = None, + auto_detect_mutables: bool = True, + tensor_numel_threshold: int = 1_000_000, + tokenizer_lock: Optional[threading.Lock] = None, + wrap_scheduler: bool = True, + ): + self._base = pipeline + self.unet = getattr(pipeline, "unet", None) + self.vae = getattr(pipeline, "vae", None) + self.text_encoder = getattr(pipeline, "text_encoder", None) + self.components = getattr(pipeline, "components", None) + + if wrap_scheduler and hasattr(pipeline, "scheduler") and pipeline.scheduler is not None: + if not isinstance(pipeline.scheduler, BaseAsyncScheduler): + pipeline.scheduler = BaseAsyncScheduler(pipeline.scheduler) + + self._mutable_attrs = list(mutable_attrs) if mutable_attrs is not None else list(self.DEFAULT_MUTABLE_ATTRS) + self._tokenizer_lock = tokenizer_lock if tokenizer_lock is not None else threading.Lock() + + self._auto_detect_mutables = bool(auto_detect_mutables) + self._tensor_numel_threshold = int(tensor_numel_threshold) + + self._auto_detected_attrs: List[str] = [] + + def _make_local_scheduler(self, num_inference_steps: int, device: Optional[str] = None, **clone_kwargs): + base_sched = getattr(self._base, "scheduler", None) + if base_sched is None: + return None + + if not isinstance(base_sched, BaseAsyncScheduler): + wrapped_scheduler = BaseAsyncScheduler(base_sched) + else: + wrapped_scheduler = base_sched + + try: + return wrapped_scheduler.clone_for_request( + num_inference_steps=num_inference_steps, device=device, **clone_kwargs + ) + except Exception as e: + logger.debug(f"clone_for_request failed: {e}; falling back to deepcopy()") + try: + return copy.deepcopy(wrapped_scheduler) + except Exception as e: + logger.warning(f"Deepcopy of scheduler failed: {e}. Returning original scheduler (*risky*).") + return wrapped_scheduler + + def _autodetect_mutables(self, max_attrs: int = 40): + if not self._auto_detect_mutables: + return [] + + if self._auto_detected_attrs: + return self._auto_detected_attrs + + candidates: List[str] = [] + seen = set() + for name in dir(self._base): + if name.startswith("__"): + continue + if name in self._mutable_attrs: + continue + if name in ("to", "save_pretrained", "from_pretrained"): + continue + try: + val = getattr(self._base, name) + except Exception: + continue + + import types + + # skip callables and modules + if callable(val) or isinstance(val, (types.ModuleType, types.FunctionType, types.MethodType)): + continue + + # containers -> candidate + if isinstance(val, (dict, list, set, tuple, bytearray)): + candidates.append(name) + seen.add(name) + else: + # try Tensor detection + try: + if isinstance(val, torch.Tensor): + if val.numel() <= self._tensor_numel_threshold: + candidates.append(name) + seen.add(name) + else: + logger.debug(f"Ignoring large tensor attr '{name}', numel={val.numel()}") + except Exception: + continue + + if len(candidates) >= max_attrs: + break + + self._auto_detected_attrs = candidates + logger.debug(f"Autodetected mutable attrs to clone: {self._auto_detected_attrs}") + return self._auto_detected_attrs + + def _is_readonly_property(self, base_obj, attr_name: str) -> bool: + try: + cls = type(base_obj) + descriptor = getattr(cls, attr_name, None) + if isinstance(descriptor, property): + return descriptor.fset is None + if hasattr(descriptor, "__set__") is False and descriptor is not None: + return False + except Exception: + pass + return False + + def _clone_mutable_attrs(self, base, local): + attrs_to_clone = list(self._mutable_attrs) + attrs_to_clone.extend(self._autodetect_mutables()) + + EXCLUDE_ATTRS = { + "components", + } + + for attr in attrs_to_clone: + if attr in EXCLUDE_ATTRS: + logger.debug(f"Skipping excluded attr '{attr}'") + continue + if not hasattr(base, attr): + continue + if self._is_readonly_property(base, attr): + logger.debug(f"Skipping read-only property '{attr}'") + continue + + try: + val = getattr(base, attr) + except Exception as e: + logger.debug(f"Could not getattr('{attr}') on base pipeline: {e}") + continue + + try: + if isinstance(val, dict): + setattr(local, attr, dict(val)) + elif isinstance(val, (list, tuple, set)): + setattr(local, attr, list(val)) + elif isinstance(val, bytearray): + setattr(local, attr, bytearray(val)) + else: + # small tensors or atomic values + if isinstance(val, torch.Tensor): + if val.numel() <= self._tensor_numel_threshold: + setattr(local, attr, val.clone()) + else: + # don't clone big tensors, keep reference + setattr(local, attr, val) + else: + try: + setattr(local, attr, copy.copy(val)) + except Exception: + setattr(local, attr, val) + except (AttributeError, TypeError) as e: + logger.debug(f"Skipping cloning attribute '{attr}' because it is not settable: {e}") + continue + except Exception as e: + logger.debug(f"Unexpected error cloning attribute '{attr}': {e}") + continue + + def _is_tokenizer_component(self, component) -> bool: + if component is None: + return False + + tokenizer_methods = ["encode", "decode", "tokenize", "__call__"] + has_tokenizer_methods = any(hasattr(component, method) for method in tokenizer_methods) + + class_name = component.__class__.__name__.lower() + has_tokenizer_in_name = "tokenizer" in class_name + + tokenizer_attrs = ["vocab_size", "pad_token", "eos_token", "bos_token"] + has_tokenizer_attrs = any(hasattr(component, attr) for attr in tokenizer_attrs) + + return has_tokenizer_methods and (has_tokenizer_in_name or has_tokenizer_attrs) + + def generate(self, *args, num_inference_steps: int = 50, device: Optional[str] = None, **kwargs): + local_scheduler = self._make_local_scheduler(num_inference_steps=num_inference_steps, device=device) + + try: + local_pipe = copy.copy(self._base) + except Exception as e: + logger.warning(f"copy.copy(self._base) failed: {e}. Falling back to deepcopy (may increase memory).") + local_pipe = copy.deepcopy(self._base) + + if local_scheduler is not None: + try: + timesteps, num_steps, configured_scheduler = async_retrieve_timesteps( + local_scheduler.scheduler, + num_inference_steps=num_inference_steps, + device=device, + return_scheduler=True, + **{k: v for k, v in kwargs.items() if k in ["timesteps", "sigmas"]}, + ) + + final_scheduler = BaseAsyncScheduler(configured_scheduler) + setattr(local_pipe, "scheduler", final_scheduler) + except Exception: + logger.warning("Could not set scheduler on local pipe; proceeding without replacing scheduler.") + + self._clone_mutable_attrs(self._base, local_pipe) + + # 4) wrap tokenizers on the local pipe with the lock wrapper + tokenizer_wrappers = {} # name -> original_tokenizer + try: + # a) wrap direct tokenizer attributes (tokenizer, tokenizer_2, ...) + for name in dir(local_pipe): + if "tokenizer" in name and not name.startswith("_"): + tok = getattr(local_pipe, name, None) + if tok is not None and self._is_tokenizer_component(tok): + tokenizer_wrappers[name] = tok + setattr( + local_pipe, + name, + lambda *args, tok=tok, **kwargs: safe_tokenize( + tok, *args, lock=self._tokenizer_lock, **kwargs + ), + ) + + # b) wrap tokenizers in components dict + if hasattr(local_pipe, "components") and isinstance(local_pipe.components, dict): + for key, val in local_pipe.components.items(): + if val is None: + continue + + if self._is_tokenizer_component(val): + tokenizer_wrappers[f"components[{key}]"] = val + local_pipe.components[key] = lambda *args, tokenizer=val, **kwargs: safe_tokenize( + tokenizer, *args, lock=self._tokenizer_lock, **kwargs + ) + + except Exception as e: + logger.debug(f"Tokenizer wrapping step encountered an error: {e}") + + result = None + cm = getattr(local_pipe, "model_cpu_offload_context", None) + try: + if callable(cm): + try: + with cm(): + result = local_pipe(*args, num_inference_steps=num_inference_steps, **kwargs) + except TypeError: + # cm might be a context manager instance rather than callable + try: + with cm: + result = local_pipe(*args, num_inference_steps=num_inference_steps, **kwargs) + except Exception as e: + logger.debug(f"model_cpu_offload_context usage failed: {e}. Proceeding without it.") + result = local_pipe(*args, num_inference_steps=num_inference_steps, **kwargs) + else: + # no offload context available — call directly + result = local_pipe(*args, num_inference_steps=num_inference_steps, **kwargs) + + return result + + finally: + try: + for name, tok in tokenizer_wrappers.items(): + if name.startswith("components["): + key = name[len("components[") : -1] + local_pipe.components[key] = tok + else: + setattr(local_pipe, name, tok) + except Exception as e: + logger.debug(f"Error restoring wrapped tokenizers: {e}") diff --git a/examples/server-async/utils/scheduler.py b/examples/server-async/utils/scheduler.py new file mode 100644 index 0000000000..86d47cac61 --- /dev/null +++ b/examples/server-async/utils/scheduler.py @@ -0,0 +1,141 @@ +import copy +import inspect +from typing import Any, List, Optional, Union + +import torch + + +class BaseAsyncScheduler: + def __init__(self, scheduler: Any): + self.scheduler = scheduler + + def __getattr__(self, name: str): + if hasattr(self.scheduler, name): + return getattr(self.scheduler, name) + raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'") + + def __setattr__(self, name: str, value): + if name == "scheduler": + super().__setattr__(name, value) + else: + if hasattr(self, "scheduler") and hasattr(self.scheduler, name): + setattr(self.scheduler, name, value) + else: + super().__setattr__(name, value) + + def clone_for_request(self, num_inference_steps: int, device: Union[str, torch.device, None] = None, **kwargs): + local = copy.deepcopy(self.scheduler) + local.set_timesteps(num_inference_steps=num_inference_steps, device=device, **kwargs) + cloned = self.__class__(local) + return cloned + + def __repr__(self): + return f"BaseAsyncScheduler({repr(self.scheduler)})" + + def __str__(self): + return f"BaseAsyncScheduler wrapping: {str(self.scheduler)}" + + +def async_retrieve_timesteps( + scheduler, + num_inference_steps: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + timesteps: Optional[List[int]] = None, + sigmas: Optional[List[float]] = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. + Handles custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Backwards compatible: by default the function behaves exactly as before and returns + (timesteps_tensor, num_inference_steps) + + If the caller passes `return_scheduler=True` in kwargs, the function will **not** mutate the passed + scheduler. Instead it will use a cloned scheduler if available (via `scheduler.clone_for_request`) + or a deepcopy fallback, call `set_timesteps` on that cloned scheduler, and return: + (timesteps_tensor, num_inference_steps, scheduler_in_use) + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`List[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`List[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Optional kwargs: + return_scheduler (bool, default False): if True, return (timesteps, num_inference_steps, scheduler_in_use) + where `scheduler_in_use` is a scheduler instance that already has timesteps set. + This mode will prefer `scheduler.clone_for_request(...)` if available, to avoid mutating the original scheduler. + + Returns: + `(timesteps_tensor, num_inference_steps)` by default (backwards compatible), or + `(timesteps_tensor, num_inference_steps, scheduler_in_use)` if `return_scheduler=True`. + """ + # pop our optional control kwarg (keeps compatibility) + return_scheduler = bool(kwargs.pop("return_scheduler", False)) + + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + + # choose scheduler to call set_timesteps on + scheduler_in_use = scheduler + if return_scheduler: + # Do not mutate the provided scheduler: prefer to clone if possible + if hasattr(scheduler, "clone_for_request"): + try: + # clone_for_request may accept num_inference_steps or other kwargs; be permissive + scheduler_in_use = scheduler.clone_for_request( + num_inference_steps=num_inference_steps or 0, device=device + ) + except Exception: + scheduler_in_use = copy.deepcopy(scheduler) + else: + # fallback deepcopy (scheduler tends to be smallish - acceptable) + scheduler_in_use = copy.deepcopy(scheduler) + + # helper to test if set_timesteps supports a particular kwarg + def _accepts(param_name: str) -> bool: + try: + return param_name in set(inspect.signature(scheduler_in_use.set_timesteps).parameters.keys()) + except (ValueError, TypeError): + # if signature introspection fails, be permissive and attempt the call later + return False + + # now call set_timesteps on the chosen scheduler_in_use (may be original or clone) + if timesteps is not None: + accepts_timesteps = _accepts("timesteps") + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler_in_use.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler_in_use.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps_out = scheduler_in_use.timesteps + num_inference_steps = len(timesteps_out) + elif sigmas is not None: + accept_sigmas = _accepts("sigmas") + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler_in_use.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler_in_use.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps_out = scheduler_in_use.timesteps + num_inference_steps = len(timesteps_out) + else: + # default path + scheduler_in_use.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps_out = scheduler_in_use.timesteps + + if return_scheduler: + return timesteps_out, num_inference_steps, scheduler_in_use + return timesteps_out, num_inference_steps diff --git a/examples/server-async/utils/utils.py b/examples/server-async/utils/utils.py new file mode 100644 index 0000000000..9f94330512 --- /dev/null +++ b/examples/server-async/utils/utils.py @@ -0,0 +1,48 @@ +import gc +import logging +import os +import tempfile +import uuid + +import torch + + +logger = logging.getLogger(__name__) + + +class Utils: + def __init__(self, host: str = "0.0.0.0", port: int = 8500): + self.service_url = f"http://{host}:{port}" + self.image_dir = os.path.join(tempfile.gettempdir(), "images") + if not os.path.exists(self.image_dir): + os.makedirs(self.image_dir) + + self.video_dir = os.path.join(tempfile.gettempdir(), "videos") + if not os.path.exists(self.video_dir): + os.makedirs(self.video_dir) + + def save_image(self, image): + if hasattr(image, "to"): + try: + image = image.to("cpu") + except Exception: + pass + + if isinstance(image, torch.Tensor): + from torchvision import transforms + + to_pil = transforms.ToPILImage() + image = to_pil(image.squeeze(0).clamp(0, 1)) + + filename = "img" + str(uuid.uuid4()).split("-")[0] + ".png" + image_path = os.path.join(self.image_dir, filename) + logger.info(f"Saving image to {image_path}") + + image.save(image_path, format="PNG", optimize=True) + + del image + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + return os.path.join(self.service_url, "images", filename) From 7e7e62c6fff46471b569cdb07b690c70514c71fc Mon Sep 17 00:00:00 2001 From: Dave Lage Date: Thu, 18 Sep 2025 03:11:22 -0400 Subject: [PATCH 006/123] Convert alphas for embedders for sd-scripts to ai toolkit conversion (#12332) * Convert alphas for embedders for sd-scripts to ai toolkit conversion * Add kohya embedders conversion test * Apply style fixes --------- Co-authored-by: Sayak Paul Co-authored-by: github-actions[bot] --- .../loaders/lora_conversion_utils.py | 86 +++++++++---------- tests/lora/test_lora_layers_flux.py | 7 ++ 2 files changed, 46 insertions(+), 47 deletions(-) diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index 6f584a5f0e..89afb6529a 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -558,70 +558,62 @@ def _convert_kohya_flux_lora_to_diffusers(state_dict): ait_sd[target_key] = value if any("guidance_in" in k for k in sds_sd): - assign_remaining_weights( - [ - ( - "time_text_embed.guidance_embedder.linear_1.{lora_key}.weight", - "lora_unet_guidance_in_in_layer.{orig_lora_key}.weight", - None, - ), - ( - "time_text_embed.guidance_embedder.linear_2.{lora_key}.weight", - "lora_unet_guidance_in_out_layer.{orig_lora_key}.weight", - None, - ), - ], + _convert_to_ai_toolkit( sds_sd, + ait_sd, + "lora_unet_guidance_in_in_layer", + "time_text_embed.guidance_embedder.linear_1", + ) + + _convert_to_ai_toolkit( + sds_sd, + ait_sd, + "lora_unet_guidance_in_out_layer", + "time_text_embed.guidance_embedder.linear_2", ) if any("img_in" in k for k in sds_sd): - assign_remaining_weights( - [ - ("x_embedder.{lora_key}.weight", "lora_unet_img_in.{orig_lora_key}.weight", None), - ], + _convert_to_ai_toolkit( sds_sd, + ait_sd, + "lora_unet_img_in", + "x_embedder", ) if any("txt_in" in k for k in sds_sd): - assign_remaining_weights( - [ - ("context_embedder.{lora_key}.weight", "lora_unet_txt_in.{orig_lora_key}.weight", None), - ], + _convert_to_ai_toolkit( sds_sd, + ait_sd, + "lora_unet_txt_in", + "context_embedder", ) if any("time_in" in k for k in sds_sd): - assign_remaining_weights( - [ - ( - "time_text_embed.timestep_embedder.linear_1.{lora_key}.weight", - "lora_unet_time_in_in_layer.{orig_lora_key}.weight", - None, - ), - ( - "time_text_embed.timestep_embedder.linear_2.{lora_key}.weight", - "lora_unet_time_in_out_layer.{orig_lora_key}.weight", - None, - ), - ], + _convert_to_ai_toolkit( sds_sd, + ait_sd, + "lora_unet_time_in_in_layer", + "time_text_embed.timestep_embedder.linear_1", + ) + _convert_to_ai_toolkit( + sds_sd, + ait_sd, + "lora_unet_time_in_out_layer", + "time_text_embed.timestep_embedder.linear_2", ) if any("vector_in" in k for k in sds_sd): - assign_remaining_weights( - [ - ( - "time_text_embed.text_embedder.linear_1.{lora_key}.weight", - "lora_unet_vector_in_in_layer.{orig_lora_key}.weight", - None, - ), - ( - "time_text_embed.text_embedder.linear_2.{lora_key}.weight", - "lora_unet_vector_in_out_layer.{orig_lora_key}.weight", - None, - ), - ], + _convert_to_ai_toolkit( sds_sd, + ait_sd, + "lora_unet_vector_in_in_layer", + "time_text_embed.text_embedder.linear_1", + ) + _convert_to_ai_toolkit( + sds_sd, + ait_sd, + "lora_unet_vector_in_out_layer", + "time_text_embed.text_embedder.linear_2", ) if any("final_layer" in k for k in sds_sd): diff --git a/tests/lora/test_lora_layers_flux.py b/tests/lora/test_lora_layers_flux.py index 7d99bcad80..e6048f509f 100644 --- a/tests/lora/test_lora_layers_flux.py +++ b/tests/lora/test_lora_layers_flux.py @@ -907,6 +907,13 @@ class FluxLoRAIntegrationTests(unittest.TestCase): assert max_diff < 1e-3 + def test_flux_kohya_embedders_conversion(self): + """Test that embedders load without throwing errors""" + self.pipeline.load_lora_weights("rockerBOO/flux-bpo-po-lora") + self.pipeline.unload_lora_weights() + + assert True + def test_flux_xlabs(self): self.pipeline.load_lora_weights("XLabs-AI/flux-lora-collection", weight_name="disney_lora.safetensors") self.pipeline.fuse_lora() From edd614ea38a7136cf66e1cb4dad696233c17358c Mon Sep 17 00:00:00 2001 From: Dhruv Nair Date: Sat, 20 Sep 2025 05:31:40 +0200 Subject: [PATCH 007/123] [CI] Fix TRANSFORMERS_FLAX_WEIGHTS_NAME import issue (#12354) update --- .../pipelines/pipeline_loading_utils.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/src/diffusers/pipelines/pipeline_loading_utils.py b/src/diffusers/pipelines/pipeline_loading_utils.py index ee767eddcc..388128df0e 100644 --- a/src/diffusers/pipelines/pipeline_loading_utils.py +++ b/src/diffusers/pipelines/pipeline_loading_utils.py @@ -48,10 +48,12 @@ from .transformers_loading_utils import _load_tokenizer_from_dduf, _load_transfo if is_transformers_available(): import transformers from transformers import PreTrainedModel, PreTrainedTokenizerBase - from transformers.utils import FLAX_WEIGHTS_NAME as TRANSFORMERS_FLAX_WEIGHTS_NAME from transformers.utils import SAFE_WEIGHTS_NAME as TRANSFORMERS_SAFE_WEIGHTS_NAME from transformers.utils import WEIGHTS_NAME as TRANSFORMERS_WEIGHTS_NAME + if is_transformers_version("<=", "4.56.2"): + from transformers.utils import FLAX_WEIGHTS_NAME as TRANSFORMERS_FLAX_WEIGHTS_NAME + if is_accelerate_available(): import accelerate from accelerate import dispatch_model @@ -112,7 +114,9 @@ def is_safetensors_compatible(filenames, passed_components=None, folder_names=No ] if is_transformers_available(): - weight_names += [TRANSFORMERS_WEIGHTS_NAME, TRANSFORMERS_SAFE_WEIGHTS_NAME, TRANSFORMERS_FLAX_WEIGHTS_NAME] + weight_names += [TRANSFORMERS_WEIGHTS_NAME, TRANSFORMERS_SAFE_WEIGHTS_NAME] + if is_transformers_version("<=", "4.56.2"): + weight_names += [TRANSFORMERS_FLAX_WEIGHTS_NAME] # model_pytorch, diffusion_model_pytorch, ... weight_prefixes = [w.split(".")[0] for w in weight_names] @@ -191,7 +195,9 @@ def filter_model_files(filenames): ] if is_transformers_available(): - weight_names += [TRANSFORMERS_WEIGHTS_NAME, TRANSFORMERS_SAFE_WEIGHTS_NAME, TRANSFORMERS_FLAX_WEIGHTS_NAME] + weight_names += [TRANSFORMERS_WEIGHTS_NAME, TRANSFORMERS_SAFE_WEIGHTS_NAME] + if is_transformers_version("<=", "4.56.2"): + weight_names += [TRANSFORMERS_FLAX_WEIGHTS_NAME] allowed_extensions = [wn.split(".")[-1] for wn in weight_names] @@ -212,7 +218,9 @@ def variant_compatible_siblings(filenames, variant=None, ignore_patterns=None) - ] if is_transformers_available(): - weight_names += [TRANSFORMERS_WEIGHTS_NAME, TRANSFORMERS_SAFE_WEIGHTS_NAME, TRANSFORMERS_FLAX_WEIGHTS_NAME] + weight_names += [TRANSFORMERS_WEIGHTS_NAME, TRANSFORMERS_SAFE_WEIGHTS_NAME] + if is_transformers_version("<=", "4.56.2"): + weight_names += [TRANSFORMERS_FLAX_WEIGHTS_NAME] # model_pytorch, diffusion_model_pytorch, ... weight_prefixes = [w.split(".")[0] for w in weight_names] From df267ee4e8500a2ef5960879f6d1ea49cc8ec40d Mon Sep 17 00:00:00 2001 From: naykun Date: Mon, 22 Sep 2025 00:10:52 +0800 Subject: [PATCH 008/123] feat: Add QwenImageEditPlus to support future feature upgrades (#12357) * feat: add support of qwenimageeditplus * add copies statement * fix copies statement * remove vl_processor reference --- src/diffusers/__init__.py | 2 + src/diffusers/pipelines/__init__.py | 2 + src/diffusers/pipelines/qwenimage/__init__.py | 2 + .../qwenimage/pipeline_qwenimage_edit.py | 1 - .../qwenimage/pipeline_qwenimage_edit_plus.py | 883 ++++++++++++++++++ .../dummy_torch_and_transformers_objects.py | 15 + 6 files changed, 904 insertions(+), 1 deletion(-) create mode 100644 src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_plus.py diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 167d39c6e8..741fcd14f2 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -515,6 +515,7 @@ else: "QwenImageControlNetPipeline", "QwenImageEditInpaintPipeline", "QwenImageEditPipeline", + "QwenImageEditPlusPipeline", "QwenImageImg2ImgPipeline", "QwenImageInpaintPipeline", "QwenImagePipeline", @@ -1170,6 +1171,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: QwenImageControlNetPipeline, QwenImageEditInpaintPipeline, QwenImageEditPipeline, + QwenImageEditPlusPipeline, QwenImageImg2ImgPipeline, QwenImageInpaintPipeline, QwenImagePipeline, diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index 17f3fc909e..190c7871d2 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -394,6 +394,7 @@ else: "QwenImageImg2ImgPipeline", "QwenImageInpaintPipeline", "QwenImageEditPipeline", + "QwenImageEditPlusPipeline", "QwenImageEditInpaintPipeline", "QwenImageControlNetInpaintPipeline", "QwenImageControlNetPipeline", @@ -721,6 +722,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: QwenImageControlNetPipeline, QwenImageEditInpaintPipeline, QwenImageEditPipeline, + QwenImageEditPlusPipeline, QwenImageImg2ImgPipeline, QwenImageInpaintPipeline, QwenImagePipeline, diff --git a/src/diffusers/pipelines/qwenimage/__init__.py b/src/diffusers/pipelines/qwenimage/__init__.py index 36d92917fd..0cd9ab40e8 100644 --- a/src/diffusers/pipelines/qwenimage/__init__.py +++ b/src/diffusers/pipelines/qwenimage/__init__.py @@ -27,6 +27,7 @@ else: _import_structure["pipeline_qwenimage_controlnet"] = ["QwenImageControlNetPipeline"] _import_structure["pipeline_qwenimage_controlnet_inpaint"] = ["QwenImageControlNetInpaintPipeline"] _import_structure["pipeline_qwenimage_edit"] = ["QwenImageEditPipeline"] + _import_structure["pipeline_qwenimage_edit_plus"] = ["QwenImageEditPlusPipeline"] _import_structure["pipeline_qwenimage_edit_inpaint"] = ["QwenImageEditInpaintPipeline"] _import_structure["pipeline_qwenimage_img2img"] = ["QwenImageImg2ImgPipeline"] _import_structure["pipeline_qwenimage_inpaint"] = ["QwenImageInpaintPipeline"] @@ -43,6 +44,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: from .pipeline_qwenimage_controlnet_inpaint import QwenImageControlNetInpaintPipeline from .pipeline_qwenimage_edit import QwenImageEditPipeline from .pipeline_qwenimage_edit_inpaint import QwenImageEditInpaintPipeline + from .pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline from .pipeline_qwenimage_img2img import QwenImageImg2ImgPipeline from .pipeline_qwenimage_inpaint import QwenImageInpaintPipeline else: diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py index 88d1ce4a46..ed37b238c8 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py +++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py @@ -208,7 +208,6 @@ class QwenImageEditPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin): # QwenImage latents are turned into 2x2 patches and packed. This means the latent width and height has to be divisible # by the patch size. So the vae scale factor is multiplied by the patch size to account for this self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) - self.vl_processor = processor self.tokenizer_max_length = 1024 self.prompt_template_encode = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n<|im_start|>assistant\n" diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_plus.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_plus.py new file mode 100644 index 0000000000..ec203edf16 --- /dev/null +++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_plus.py @@ -0,0 +1,883 @@ +# Copyright 2025 Qwen-Image Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import math +from typing import Any, Callable, Dict, List, Optional, Union + +import numpy as np +import torch +from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessor + +from ...image_processor import PipelineImageInput, VaeImageProcessor +from ...loaders import QwenImageLoraLoaderMixin +from ...models import AutoencoderKLQwenImage, QwenImageTransformer2DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import is_torch_xla_available, logging, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ..pipeline_utils import DiffusionPipeline +from .pipeline_output import QwenImagePipelineOutput + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from PIL import Image + >>> from diffusers import QwenImageEditPlusPipeline + >>> from diffusers.utils import load_image + + >>> pipe = QwenImageEditPlusPipeline.from_pretrained("Qwen/Qwen-Image-Edit-2509", torch_dtype=torch.bfloat16) + >>> pipe.to("cuda") + >>> image = load_image( + ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/yarn-art-pikachu.png" + ... ).convert("RGB") + >>> prompt = ( + ... "Make Pikachu hold a sign that says 'Qwen Edit is awesome', yarn art style, detailed, vibrant colors" + ... ) + >>> # Depending on the variant being used, the pipeline call will slightly vary. + >>> # Refer to the pipeline documentation for more details. + >>> image = pipe(image, prompt, num_inference_steps=50).images[0] + >>> image.save("qwenimage_edit_plus.png") + ``` +""" + +CONDITION_IMAGE_SIZE = 384 * 384 +VAE_IMAGE_SIZE = 1024 * 1024 + + +# Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.calculate_shift +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + timesteps: Optional[List[int]] = None, + sigmas: Optional[List[float]] = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`List[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`List[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents +def retrieve_latents( + encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" +): + if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": + return encoder_output.latent_dist.sample(generator) + elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": + return encoder_output.latent_dist.mode() + elif hasattr(encoder_output, "latents"): + return encoder_output.latents + else: + raise AttributeError("Could not access latents of provided encoder_output") + + +def calculate_dimensions(target_area, ratio): + width = math.sqrt(target_area * ratio) + height = width / ratio + + width = round(width / 32) * 32 + height = round(height / 32) * 32 + + return width, height + + +class QwenImageEditPlusPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin): + r""" + The Qwen-Image-Edit pipeline for image editing. + + Args: + transformer ([`QwenImageTransformer2DModel`]): + Conditional Transformer (MMDiT) architecture to denoise the encoded image latents. + scheduler ([`FlowMatchEulerDiscreteScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded image latents. + vae ([`AutoencoderKL`]): + Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. + text_encoder ([`Qwen2.5-VL-7B-Instruct`]): + [Qwen2.5-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct), specifically the + [Qwen2.5-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct) variant. + tokenizer (`QwenTokenizer`): + Tokenizer of class + [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer). + """ + + model_cpu_offload_seq = "text_encoder->transformer->vae" + _callback_tensor_inputs = ["latents", "prompt_embeds"] + + def __init__( + self, + scheduler: FlowMatchEulerDiscreteScheduler, + vae: AutoencoderKLQwenImage, + text_encoder: Qwen2_5_VLForConditionalGeneration, + tokenizer: Qwen2Tokenizer, + processor: Qwen2VLProcessor, + transformer: QwenImageTransformer2DModel, + ): + super().__init__() + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + processor=processor, + transformer=transformer, + scheduler=scheduler, + ) + self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8 + self.latent_channels = self.vae.config.z_dim if getattr(self, "vae", None) else 16 + # QwenImage latents are turned into 2x2 patches and packed. This means the latent width and height has to be divisible + # by the patch size. So the vae scale factor is multiplied by the patch size to account for this + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) + self.tokenizer_max_length = 1024 + + self.prompt_template_encode = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" + self.prompt_template_encode_start_idx = 64 + self.default_sample_size = 128 + + # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._extract_masked_hidden + def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor): + bool_mask = mask.bool() + valid_lengths = bool_mask.sum(dim=1) + selected = hidden_states[bool_mask] + split_result = torch.split(selected, valid_lengths.tolist(), dim=0) + + return split_result + + def _get_qwen_prompt_embeds( + self, + prompt: Union[str, List[str]] = None, + image: Optional[torch.Tensor] = None, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + img_prompt_template = "Picture {}: <|vision_start|><|image_pad|><|vision_end|>" + if isinstance(image, list): + base_img_prompt = "" + for i, img in enumerate(image): + base_img_prompt += img_prompt_template.format(i + 1) + elif image is not None: + base_img_prompt = img_prompt_template.format(1) + else: + base_img_prompt = "" + + template = self.prompt_template_encode + + drop_idx = self.prompt_template_encode_start_idx + txt = [template.format(base_img_prompt + e) for e in prompt] + + model_inputs = self.processor( + text=txt, + images=image, + padding=True, + return_tensors="pt", + ).to(device) + + outputs = self.text_encoder( + input_ids=model_inputs.input_ids, + attention_mask=model_inputs.attention_mask, + pixel_values=model_inputs.pixel_values, + image_grid_thw=model_inputs.image_grid_thw, + output_hidden_states=True, + ) + + hidden_states = outputs.hidden_states[-1] + split_hidden_states = self._extract_masked_hidden(hidden_states, model_inputs.attention_mask) + split_hidden_states = [e[drop_idx:] for e in split_hidden_states] + attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states] + max_seq_len = max([e.size(0) for e in split_hidden_states]) + prompt_embeds = torch.stack( + [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states] + ) + encoder_attention_mask = torch.stack( + [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list] + ) + + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + + return prompt_embeds, encoder_attention_mask + + # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline.encode_prompt + def encode_prompt( + self, + prompt: Union[str, List[str]], + image: Optional[torch.Tensor] = None, + device: Optional[torch.device] = None, + num_images_per_prompt: int = 1, + prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds_mask: Optional[torch.Tensor] = None, + max_sequence_length: int = 1024, + ): + r""" + + Args: + prompt (`str` or `List[str]`, *optional*): + prompt to be encoded + image (`torch.Tensor`, *optional*): + image to be encoded + device: (`torch.device`): + torch device + num_images_per_prompt (`int`): + number of images that should be generated per prompt + 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. + """ + device = device or self._execution_device + + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0] + + if prompt_embeds is None: + prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, image, device) + + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1) + prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len) + + return prompt_embeds, prompt_embeds_mask + + # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline.check_inputs + def check_inputs( + self, + prompt, + height, + width, + negative_prompt=None, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_embeds_mask=None, + negative_prompt_embeds_mask=None, + callback_on_step_end_tensor_inputs=None, + max_sequence_length=None, + ): + if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0: + logger.warning( + f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} but are {height} and {width}. Dimensions will be resized accordingly" + ) + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + if prompt_embeds is not None and prompt_embeds_mask is None: + raise ValueError( + "If `prompt_embeds` are provided, `prompt_embeds_mask` also have to be passed. Make sure to generate `prompt_embeds_mask` from the same text encoder that was used to generate `prompt_embeds`." + ) + if negative_prompt_embeds is not None and negative_prompt_embeds_mask is None: + raise ValueError( + "If `negative_prompt_embeds` are provided, `negative_prompt_embeds_mask` also have to be passed. Make sure to generate `negative_prompt_embeds_mask` from the same text encoder that was used to generate `negative_prompt_embeds`." + ) + + if max_sequence_length is not None and max_sequence_length > 1024: + raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}") + + @staticmethod + # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._pack_latents + def _pack_latents(latents, batch_size, num_channels_latents, height, width): + latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2) + latents = latents.permute(0, 2, 4, 1, 3, 5) + latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4) + + return latents + + @staticmethod + # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._unpack_latents + def _unpack_latents(latents, height, width, vae_scale_factor): + batch_size, num_patches, channels = latents.shape + + # VAE applies 8x compression on images but we must also account for packing which requires + # latent height and width to be divisible by 2. + height = 2 * (int(height) // (vae_scale_factor * 2)) + width = 2 * (int(width) // (vae_scale_factor * 2)) + + latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2) + latents = latents.permute(0, 3, 1, 4, 2, 5) + + latents = latents.reshape(batch_size, channels // (2 * 2), 1, height, width) + + return latents + + # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline._encode_vae_image + def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): + if isinstance(generator, list): + image_latents = [ + retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i], sample_mode="argmax") + 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, sample_mode="argmax") + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.latent_channels, 1, 1, 1) + .to(image_latents.device, image_latents.dtype) + ) + latents_std = ( + torch.tensor(self.vae.config.latents_std) + .view(1, self.latent_channels, 1, 1, 1) + .to(image_latents.device, image_latents.dtype) + ) + image_latents = (image_latents - latents_mean) / latents_std + + return image_latents + + def prepare_latents( + self, + images, + batch_size, + num_channels_latents, + height, + width, + dtype, + device, + generator, + latents=None, + ): + # VAE applies 8x compression on images but we must also account for packing which requires + # latent height and width to be divisible by 2. + height = 2 * (int(height) // (self.vae_scale_factor * 2)) + width = 2 * (int(width) // (self.vae_scale_factor * 2)) + + shape = (batch_size, 1, num_channels_latents, height, width) + + image_latents = None + if images is not None: + if not isinstance(images, list): + images = [images] + all_image_latents = [] + for image in images: + image = image.to(device=device, dtype=dtype) + if image.shape[1] != self.latent_channels: + image_latents = self._encode_vae_image(image=image, generator=generator) + else: + image_latents = image + if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0: + # expand init_latents for batch_size + additional_image_per_prompt = batch_size // image_latents.shape[0] + image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0) + elif batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0: + raise ValueError( + f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts." + ) + else: + image_latents = torch.cat([image_latents], dim=0) + + image_latent_height, image_latent_width = image_latents.shape[3:] + image_latents = self._pack_latents( + image_latents, batch_size, num_channels_latents, image_latent_height, image_latent_width + ) + all_image_latents.append(image_latents) + image_latents = torch.cat(all_image_latents, dim=1) + + 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) + latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width) + else: + latents = latents.to(device=device, dtype=dtype) + + return latents, image_latents + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def attention_kwargs(self): + return self._attention_kwargs + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + image: Optional[PipelineImageInput] = None, + prompt: Union[str, List[str]] = None, + negative_prompt: Union[str, List[str]] = None, + true_cfg_scale: float = 4.0, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 50, + sigmas: Optional[List[float]] = None, + guidance_scale: Optional[float] = None, + num_images_per_prompt: int = 1, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds_mask: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt_embeds_mask: Optional[torch.Tensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + attention_kwargs: Optional[Dict[str, Any]] = None, + callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + max_sequence_length: int = 512, + ): + r""" + Function invoked when calling the pipeline for generation. + + Args: + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both + numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list + or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a + list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image + latents as `image`, but if passing latents directly it is not encoded again. + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + 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 `true_cfg_scale` is + not greater than `1`). + true_cfg_scale (`float`, *optional*, defaults to 1.0): + true_cfg_scale (`float`, *optional*, defaults to 1.0): Guidance scale as defined in [Classifier-Free + Diffusion Guidance](https://huggingface.co/papers/2207.12598). `true_cfg_scale` is defined as `w` of + equation 2. of [Imagen Paper](https://huggingface.co/papers/2205.11487). Classifier-free guidance is + enabled by setting `true_cfg_scale > 1` and a provided `negative_prompt`. Higher guidance scale + encourages to generate images that are closely linked to the text `prompt`, usually at the expense of + lower image quality. + height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The height in pixels of the generated image. This is set to 1024 by default for the best results. + width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The width in pixels of the generated image. This is set to 1024 by default for the best results. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + sigmas (`List[float]`, *optional*): + Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in + their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed + will be used. + guidance_scale (`float`, *optional*, defaults to None): + A guidance scale value for guidance distilled models. Unlike the traditional classifier-free guidance + where the guidance scale is applied during inference through noise prediction rescaling, guidance + distilled models take the guidance scale directly as an input parameter during forward pass. 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. This + parameter in the pipeline is there to support future guidance-distilled models when they come up. It is + ignored when not using guidance distilled models. To enable traditional classifier-free guidance, + please pass `true_cfg_scale > 1.0` and `negative_prompt` (even an empty negative prompt like " " should + enable classifier-free guidance computations). + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will be generated by sampling using the supplied random `generator`. + prompt_embeds (`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. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.qwenimage.QwenImagePipelineOutput`] instead of a plain tuple. + attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + callback_on_step_end (`Callable`, *optional*): + A function that calls at the end of each denoising steps during the inference. The function is called + with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, + callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by + `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`. + + Examples: + + Returns: + [`~pipelines.qwenimage.QwenImagePipelineOutput`] or `tuple`: + [`~pipelines.qwenimage.QwenImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When + returning a tuple, the first element is a list with the generated images. + """ + image_size = image[-1].size if isinstance(image, list) else image.size + calculated_width, calculated_height = calculate_dimensions(1024 * 1024, image_size[0] / image_size[1]) + height = height or calculated_height + width = width or calculated_width + + multiple_of = self.vae_scale_factor * 2 + width = width // multiple_of * multiple_of + height = height // multiple_of * multiple_of + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt, + height, + width, + negative_prompt=negative_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_embeds_mask=prompt_embeds_mask, + negative_prompt_embeds_mask=negative_prompt_embeds_mask, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + max_sequence_length=max_sequence_length, + ) + + self._guidance_scale = guidance_scale + self._attention_kwargs = attention_kwargs + self._current_timestep = None + self._interrupt = False + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self._execution_device + # 3. Preprocess image + if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels): + if not isinstance(image, list): + image = [image] + condition_image_sizes = [] + condition_images = [] + vae_image_sizes = [] + vae_images = [] + for img in image: + image_width, image_height = img.size + condition_width, condition_height = calculate_dimensions( + CONDITION_IMAGE_SIZE, image_width / image_height + ) + vae_width, vae_height = calculate_dimensions(VAE_IMAGE_SIZE, image_width / image_height) + condition_image_sizes.append((condition_width, condition_height)) + vae_image_sizes.append((vae_width, vae_height)) + condition_images.append(self.image_processor.resize(img, condition_height, condition_width)) + vae_images.append(self.image_processor.preprocess(img, vae_height, vae_width).unsqueeze(2)) + + has_neg_prompt = negative_prompt is not None or ( + negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None + ) + + if true_cfg_scale > 1 and not has_neg_prompt: + logger.warning( + f"true_cfg_scale is passed as {true_cfg_scale}, but classifier-free guidance is not enabled since no negative_prompt is provided." + ) + elif true_cfg_scale <= 1 and has_neg_prompt: + logger.warning( + " negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1" + ) + + do_true_cfg = true_cfg_scale > 1 and has_neg_prompt + prompt_embeds, prompt_embeds_mask = self.encode_prompt( + image=condition_images, + prompt=prompt, + prompt_embeds=prompt_embeds, + prompt_embeds_mask=prompt_embeds_mask, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + ) + if do_true_cfg: + negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt( + image=condition_images, + prompt=negative_prompt, + prompt_embeds=negative_prompt_embeds, + prompt_embeds_mask=negative_prompt_embeds_mask, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + ) + + # 4. Prepare latent variables + num_channels_latents = self.transformer.config.in_channels // 4 + latents, image_latents = self.prepare_latents( + vae_images, + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + prompt_embeds.dtype, + device, + generator, + latents, + ) + img_shapes = [ + [ + (1, height // self.vae_scale_factor // 2, width // self.vae_scale_factor // 2), + *[ + (1, vae_height // self.vae_scale_factor // 2, vae_width // self.vae_scale_factor // 2) + for vae_width, vae_height in vae_image_sizes + ], + ] + ] * batch_size + + # 5. Prepare timesteps + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas + image_seq_len = latents.shape[1] + mu = calculate_shift( + image_seq_len, + self.scheduler.config.get("base_image_seq_len", 256), + self.scheduler.config.get("max_image_seq_len", 4096), + self.scheduler.config.get("base_shift", 0.5), + self.scheduler.config.get("max_shift", 1.15), + ) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + sigmas=sigmas, + mu=mu, + ) + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + + # handle guidance + if self.transformer.config.guidance_embeds and guidance_scale is None: + raise ValueError("guidance_scale is required for guidance-distilled model.") + elif self.transformer.config.guidance_embeds: + guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32) + guidance = guidance.expand(latents.shape[0]) + elif not self.transformer.config.guidance_embeds and guidance_scale is not None: + logger.warning( + f"guidance_scale is passed as {guidance_scale}, but ignored since the model is not guidance-distilled." + ) + guidance = None + elif not self.transformer.config.guidance_embeds and guidance_scale is None: + guidance = None + + if self.attention_kwargs is None: + self._attention_kwargs = {} + + txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None + negative_txt_seq_lens = ( + negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None + ) + + # 6. Denoising loop + self.scheduler.set_begin_index(0) + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + self._current_timestep = t + + latent_model_input = latents + if image_latents is not None: + latent_model_input = torch.cat([latents, image_latents], dim=1) + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latents.shape[0]).to(latents.dtype) + with self.transformer.cache_context("cond"): + noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep / 1000, + guidance=guidance, + encoder_hidden_states_mask=prompt_embeds_mask, + encoder_hidden_states=prompt_embeds, + img_shapes=img_shapes, + txt_seq_lens=txt_seq_lens, + attention_kwargs=self.attention_kwargs, + return_dict=False, + )[0] + noise_pred = noise_pred[:, : latents.size(1)] + + if do_true_cfg: + with self.transformer.cache_context("uncond"): + neg_noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep / 1000, + guidance=guidance, + encoder_hidden_states_mask=negative_prompt_embeds_mask, + encoder_hidden_states=negative_prompt_embeds, + img_shapes=img_shapes, + txt_seq_lens=negative_txt_seq_lens, + attention_kwargs=self.attention_kwargs, + return_dict=False, + )[0] + neg_noise_pred = neg_noise_pred[:, : latents.size(1)] + comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred) + + cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True) + noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True) + noise_pred = comb_pred * (cond_norm / noise_norm) + + # compute the previous noisy sample x_t -> x_t-1 + latents_dtype = latents.dtype + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + + if latents.dtype != latents_dtype: + if torch.backends.mps.is_available(): + # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272 + latents = latents.to(latents_dtype) + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + self._current_timestep = None + if output_type == "latent": + image = latents + else: + latents = self._unpack_latents(latents, height, width, self.vae_scale_factor) + latents = latents.to(self.vae.dtype) + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( + latents.device, latents.dtype + ) + latents = latents / latents_std + latents_mean + image = self.vae.decode(latents, return_dict=False)[0][:, :, 0] + image = self.image_processor.postprocess(image, output_type=output_type) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return QwenImagePipelineOutput(images=image) diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index e29be174f0..bb8fea8c8a 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -1892,6 +1892,21 @@ class QwenImageEditPipeline(metaclass=DummyObject): requires_backends(cls, ["torch", "transformers"]) +class QwenImageEditPlusPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class QwenImageImg2ImgPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] From c006a95df1f84f3c3c49b53926bd75c518a68fcb Mon Sep 17 00:00:00 2001 From: Jason Cox Date: Sun, 21 Sep 2025 20:07:17 -0700 Subject: [PATCH 009/123] Fix example server install instructions (#12362) * Upgrade huggingface-hub to version 0.35.0 Updated huggingface-hub version from 0.26.1 to 0.35.0. * Add uvicorn and accelerate to requirements * Fix install instructions for server --- examples/server/README.md | 4 ++-- examples/server/requirements.in | 3 ++- examples/server/requirements.txt | 2 +- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/examples/server/README.md b/examples/server/README.md index 8ad0ed3cbe..f8cd58fc1c 100644 --- a/examples/server/README.md +++ b/examples/server/README.md @@ -9,8 +9,8 @@ This guide will show you how to use the [`StableDiffusion3Pipeline`] in a server Start by navigating to the `examples/server` folder and installing all of the dependencies. ```py -pip install . -pip install -f requirements.txt +pip install diffusers +pip install -r requirements.txt ``` Launch the server with the following command. diff --git a/examples/server/requirements.in b/examples/server/requirements.in index a469569a10..f8c35d48cd 100644 --- a/examples/server/requirements.in +++ b/examples/server/requirements.in @@ -6,4 +6,5 @@ py-consul prometheus_client >= 0.18.0 prometheus-fastapi-instrumentator >= 7.0.0 fastapi -uvicorn \ No newline at end of file +uvicorn +accelerate diff --git a/examples/server/requirements.txt b/examples/server/requirements.txt index b91a8861a0..688a4ee94f 100644 --- a/examples/server/requirements.txt +++ b/examples/server/requirements.txt @@ -39,7 +39,7 @@ fsspec==2024.10.0 # torch h11==0.14.0 # via uvicorn -huggingface-hub==0.26.1 +huggingface-hub==0.35.0 # via # tokenizers # transformers From 843355f89fd043e82b3344d9259e6faa640da6f9 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Mon, 22 Sep 2025 11:17:47 +0530 Subject: [PATCH 010/123] [tests] xfail some kandinsky tests. (#12364) xfail some kandinsky tests. --- src/diffusers/pipelines/qwenimage/__init__.py | 2 +- tests/pipelines/kandinsky/test_kandinsky.py | 5 +++++ tests/pipelines/kandinsky/test_kandinsky_combined.py | 11 +++++++++++ tests/pipelines/kandinsky/test_kandinsky_img2img.py | 5 +++++ tests/pipelines/kandinsky/test_kandinsky_inpaint.py | 5 +++++ 5 files changed, 27 insertions(+), 1 deletion(-) diff --git a/src/diffusers/pipelines/qwenimage/__init__.py b/src/diffusers/pipelines/qwenimage/__init__.py index 0cd9ab40e8..2400632ba2 100644 --- a/src/diffusers/pipelines/qwenimage/__init__.py +++ b/src/diffusers/pipelines/qwenimage/__init__.py @@ -27,8 +27,8 @@ else: _import_structure["pipeline_qwenimage_controlnet"] = ["QwenImageControlNetPipeline"] _import_structure["pipeline_qwenimage_controlnet_inpaint"] = ["QwenImageControlNetInpaintPipeline"] _import_structure["pipeline_qwenimage_edit"] = ["QwenImageEditPipeline"] - _import_structure["pipeline_qwenimage_edit_plus"] = ["QwenImageEditPlusPipeline"] _import_structure["pipeline_qwenimage_edit_inpaint"] = ["QwenImageEditInpaintPipeline"] + _import_structure["pipeline_qwenimage_edit_plus"] = ["QwenImageEditPlusPipeline"] _import_structure["pipeline_qwenimage_img2img"] = ["QwenImageImg2ImgPipeline"] _import_structure["pipeline_qwenimage_inpaint"] = ["QwenImageInpaintPipeline"] diff --git a/tests/pipelines/kandinsky/test_kandinsky.py b/tests/pipelines/kandinsky/test_kandinsky.py index 911c6e49ae..9fa39b1bf5 100644 --- a/tests/pipelines/kandinsky/test_kandinsky.py +++ b/tests/pipelines/kandinsky/test_kandinsky.py @@ -18,11 +18,13 @@ import random import unittest import numpy as np +import pytest import torch from transformers import XLMRobertaTokenizerFast from diffusers import DDIMScheduler, KandinskyPipeline, KandinskyPriorPipeline, UNet2DConditionModel, VQModel from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP +from diffusers.utils import is_transformers_version from ...testing_utils import ( backend_empty_cache, @@ -215,6 +217,9 @@ class KandinskyPipelineFastTests(PipelineTesterMixin, unittest.TestCase): dummy = Dummies() return dummy.get_dummy_inputs(device=device, seed=seed) + @pytest.mark.xfail( + condition=is_transformers_version(">=", "4.56.2"), reason="Latest transformers changes the slices", strict=True + ) def test_kandinsky(self): device = "cpu" diff --git a/tests/pipelines/kandinsky/test_kandinsky_combined.py b/tests/pipelines/kandinsky/test_kandinsky_combined.py index d744d10821..ca80461d87 100644 --- a/tests/pipelines/kandinsky/test_kandinsky_combined.py +++ b/tests/pipelines/kandinsky/test_kandinsky_combined.py @@ -16,8 +16,10 @@ import unittest import numpy as np +import pytest from diffusers import KandinskyCombinedPipeline, KandinskyImg2ImgCombinedPipeline, KandinskyInpaintCombinedPipeline +from diffusers.utils import is_transformers_version from ...testing_utils import enable_full_determinism, require_torch_accelerator, torch_device from ..test_pipelines_common import PipelineTesterMixin @@ -73,6 +75,9 @@ class KandinskyPipelineCombinedFastTests(PipelineTesterMixin, unittest.TestCase) ) return inputs + @pytest.mark.xfail( + condition=is_transformers_version(">=", "4.56.2"), reason="Latest transformers changes the slices", strict=True + ) def test_kandinsky(self): device = "cpu" @@ -181,6 +186,9 @@ class KandinskyPipelineImg2ImgCombinedFastTests(PipelineTesterMixin, unittest.Te inputs.pop("negative_image_embeds") return inputs + @pytest.mark.xfail( + condition=is_transformers_version(">=", "4.56.2"), reason="Latest transformers changes the slices", strict=True + ) def test_kandinsky(self): device = "cpu" @@ -292,6 +300,9 @@ class KandinskyPipelineInpaintCombinedFastTests(PipelineTesterMixin, unittest.Te inputs.pop("negative_image_embeds") return inputs + @pytest.mark.xfail( + condition=is_transformers_version(">=", "4.56.2"), reason="Latest transformers changes the slices", strict=True + ) def test_kandinsky(self): device = "cpu" diff --git a/tests/pipelines/kandinsky/test_kandinsky_img2img.py b/tests/pipelines/kandinsky/test_kandinsky_img2img.py index 4074c8db22..6bcd9587f2 100644 --- a/tests/pipelines/kandinsky/test_kandinsky_img2img.py +++ b/tests/pipelines/kandinsky/test_kandinsky_img2img.py @@ -18,6 +18,7 @@ import random import unittest import numpy as np +import pytest import torch from PIL import Image from transformers import XLMRobertaTokenizerFast @@ -31,6 +32,7 @@ from diffusers import ( VQModel, ) from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP +from diffusers.utils import is_transformers_version from ...testing_utils import ( backend_empty_cache, @@ -237,6 +239,9 @@ class KandinskyImg2ImgPipelineFastTests(PipelineTesterMixin, unittest.TestCase): dummies = Dummies() return dummies.get_dummy_inputs(device=device, seed=seed) + @pytest.mark.xfail( + condition=is_transformers_version(">=", "4.56.2"), reason="Latest transformers changes the slices", strict=True + ) def test_kandinsky_img2img(self): device = "cpu" diff --git a/tests/pipelines/kandinsky/test_kandinsky_inpaint.py b/tests/pipelines/kandinsky/test_kandinsky_inpaint.py index b789a63cdd..6383ca71ef 100644 --- a/tests/pipelines/kandinsky/test_kandinsky_inpaint.py +++ b/tests/pipelines/kandinsky/test_kandinsky_inpaint.py @@ -18,12 +18,14 @@ import random import unittest import numpy as np +import pytest import torch from PIL import Image from transformers import XLMRobertaTokenizerFast from diffusers import DDIMScheduler, KandinskyInpaintPipeline, KandinskyPriorPipeline, UNet2DConditionModel, VQModel from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP +from diffusers.utils import is_transformers_version from ...testing_utils import ( backend_empty_cache, @@ -231,6 +233,9 @@ class KandinskyInpaintPipelineFastTests(PipelineTesterMixin, unittest.TestCase): dummies = Dummies() return dummies.get_dummy_inputs(device=device, seed=seed) + @pytest.mark.xfail( + condition=is_transformers_version(">=", "4.56.2"), reason="Latest transformers changes the slices", strict=True + ) def test_kandinsky_inpaint(self): device = "cpu" From d83d35c1bba322803c1631c6683e43f93dbdec37 Mon Sep 17 00:00:00 2001 From: Chen Mingyi <65432252+mingyi456@users.noreply.github.com> Date: Mon, 22 Sep 2025 14:55:34 +0800 Subject: [PATCH 011/123] Fix bug with VAE slicing in autoencoder_dc.py (#12343) --- src/diffusers/models/autoencoders/autoencoder_dc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_dc.py b/src/diffusers/models/autoencoders/autoencoder_dc.py index d3f31de854..783f22e97d 100644 --- a/src/diffusers/models/autoencoders/autoencoder_dc.py +++ b/src/diffusers/models/autoencoders/autoencoder_dc.py @@ -617,7 +617,7 @@ class AutoencoderDC(ModelMixin, ConfigMixin, FromOriginalModelMixin): returned. """ if self.use_slicing and z.size(0) > 1: - decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)] + decoded_slices = [self._decode(z_slice) for z_slice in z.split(1)] decoded = torch.cat(decoded_slices) else: decoded = self._decode(z) From 78031c29383852efce15d8d9a79d2d985394623a Mon Sep 17 00:00:00 2001 From: SahilCarterr <110806554+SahilCarterr@users.noreply.github.com> Date: Mon, 22 Sep 2025 12:37:41 +0530 Subject: [PATCH 012/123] [Fix] enable_xformers_memory_efficient_attention() in Flux Pipeline (#12337) * FIxes enable_xformers_memory_efficient_attention() * Update attention.py --- src/diffusers/models/attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/models/attention.py b/src/diffusers/models/attention.py index c99133f257..3f59c8da8e 100644 --- a/src/diffusers/models/attention.py +++ b/src/diffusers/models/attention.py @@ -241,7 +241,7 @@ class AttentionModuleMixin: op_fw, op_bw = attention_op dtype, *_ = op_fw.SUPPORTED_DTYPES q = torch.randn((1, 2, 40), device="cuda", dtype=dtype) - _ = xops.memory_efficient_attention(q, q, q) + _ = xops.ops.memory_efficient_attention(q, q, q) except Exception as e: raise e From d8310a8fca812ff4afdcc1fd09c135dcccaab670 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Mon, 22 Sep 2025 15:14:39 +0530 Subject: [PATCH 013/123] [lora] factor out the overlaps in `save_lora_weights()`. (#12027) * factor out the overlaps in save_lora_weights(). * remove comment. * remove comment. * up * fix-copies --- src/diffusers/loaders/lora_base.py | 35 ++ src/diffusers/loaders/lora_pipeline.py | 426 ++++++++++--------------- 2 files changed, 205 insertions(+), 256 deletions(-) diff --git a/src/diffusers/loaders/lora_base.py b/src/diffusers/loaders/lora_base.py index d18c82df4f..0ee32f820b 100644 --- a/src/diffusers/loaders/lora_base.py +++ b/src/diffusers/loaders/lora_base.py @@ -1064,6 +1064,41 @@ class LoraBaseMixin: save_function(state_dict, save_path) logger.info(f"Model weights saved in {save_path}") + @classmethod + def _save_lora_weights( + cls, + save_directory: Union[str, os.PathLike], + lora_layers: Dict[str, Dict[str, Union[torch.nn.Module, torch.Tensor]]], + lora_metadata: Dict[str, Optional[dict]], + is_main_process: bool = True, + weight_name: str = None, + save_function: Callable = None, + safe_serialization: bool = True, + ): + """ + Helper method to pack and save LoRA weights and metadata. This method centralizes the saving logic for all + pipeline types. + """ + state_dict = {} + final_lora_adapter_metadata = {} + + for prefix, layers in lora_layers.items(): + state_dict.update(cls.pack_weights(layers, prefix)) + + for prefix, metadata in lora_metadata.items(): + if metadata: + final_lora_adapter_metadata.update(_pack_dict_with_prefix(metadata, prefix)) + + cls.write_lora_layers( + state_dict=state_dict, + save_directory=save_directory, + is_main_process=is_main_process, + weight_name=weight_name, + save_function=save_function, + safe_serialization=safe_serialization, + lora_adapter_metadata=final_lora_adapter_metadata if final_lora_adapter_metadata else None, + ) + @classmethod def _optionally_disable_offloading(cls, _pipeline): return _func_optionally_disable_offloading(_pipeline=_pipeline) diff --git a/src/diffusers/loaders/lora_pipeline.py b/src/diffusers/loaders/lora_pipeline.py index 7e89066f1f..8060b519f1 100644 --- a/src/diffusers/loaders/lora_pipeline.py +++ b/src/diffusers/loaders/lora_pipeline.py @@ -510,35 +510,28 @@ class StableDiffusionLoraLoaderMixin(LoraBaseMixin): text_encoder_lora_adapter_metadata: LoRA adapter metadata associated with the text encoder to be serialized with the state dict. """ - state_dict = {} - lora_adapter_metadata = {} - - if not (unet_lora_layers or text_encoder_lora_layers): - raise ValueError("You must pass at least one of `unet_lora_layers` and `text_encoder_lora_layers`.") + lora_layers = {} + lora_metadata = {} if unet_lora_layers: - state_dict.update(cls.pack_weights(unet_lora_layers, cls.unet_name)) + lora_layers[cls.unet_name] = unet_lora_layers + lora_metadata[cls.unet_name] = unet_lora_adapter_metadata if text_encoder_lora_layers: - state_dict.update(cls.pack_weights(text_encoder_lora_layers, cls.text_encoder_name)) + lora_layers[cls.text_encoder_name] = text_encoder_lora_layers + lora_metadata[cls.text_encoder_name] = text_encoder_lora_adapter_metadata - if unet_lora_adapter_metadata: - lora_adapter_metadata.update(_pack_dict_with_prefix(unet_lora_adapter_metadata, cls.unet_name)) + if not lora_layers: + raise ValueError("You must pass at least one of `unet_lora_layers` or `text_encoder_lora_layers`.") - if text_encoder_lora_adapter_metadata: - lora_adapter_metadata.update( - _pack_dict_with_prefix(text_encoder_lora_adapter_metadata, cls.text_encoder_name) - ) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, + cls._save_lora_weights( save_directory=save_directory, + lora_layers=lora_layers, + lora_metadata=lora_metadata, is_main_process=is_main_process, weight_name=weight_name, save_function=save_function, safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, ) def fuse_lora( @@ -1004,44 +997,34 @@ class StableDiffusionXLLoraLoaderMixin(LoraBaseMixin): text_encoder_2_lora_adapter_metadata: LoRA adapter metadata associated with the second text encoder to be serialized with the state dict. """ - state_dict = {} - lora_adapter_metadata = {} - - if not (unet_lora_layers or text_encoder_lora_layers or text_encoder_2_lora_layers): - raise ValueError( - "You must pass at least one of `unet_lora_layers`, `text_encoder_lora_layers`, `text_encoder_2_lora_layers`." - ) + lora_layers = {} + lora_metadata = {} if unet_lora_layers: - state_dict.update(cls.pack_weights(unet_lora_layers, cls.unet_name)) + lora_layers[cls.unet_name] = unet_lora_layers + lora_metadata[cls.unet_name] = unet_lora_adapter_metadata if text_encoder_lora_layers: - state_dict.update(cls.pack_weights(text_encoder_lora_layers, "text_encoder")) + lora_layers["text_encoder"] = text_encoder_lora_layers + lora_metadata["text_encoder"] = text_encoder_lora_adapter_metadata if text_encoder_2_lora_layers: - state_dict.update(cls.pack_weights(text_encoder_2_lora_layers, "text_encoder_2")) + lora_layers["text_encoder_2"] = text_encoder_2_lora_layers + lora_metadata["text_encoder_2"] = text_encoder_2_lora_adapter_metadata - if unet_lora_adapter_metadata is not None: - lora_adapter_metadata.update(_pack_dict_with_prefix(unet_lora_adapter_metadata, cls.unet_name)) - - if text_encoder_lora_adapter_metadata: - lora_adapter_metadata.update( - _pack_dict_with_prefix(text_encoder_lora_adapter_metadata, cls.text_encoder_name) + if not lora_layers: + raise ValueError( + "You must pass at least one of `unet_lora_layers`, `text_encoder_lora_layers`, or `text_encoder_2_lora_layers`." ) - if text_encoder_2_lora_adapter_metadata: - lora_adapter_metadata.update( - _pack_dict_with_prefix(text_encoder_2_lora_adapter_metadata, "text_encoder_2") - ) - - cls.write_lora_layers( - state_dict=state_dict, + cls._save_lora_weights( save_directory=save_directory, + lora_layers=lora_layers, + lora_metadata=lora_metadata, is_main_process=is_main_process, weight_name=weight_name, save_function=save_function, safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, ) def fuse_lora( @@ -1467,46 +1450,34 @@ class SD3LoraLoaderMixin(LoraBaseMixin): text_encoder_2_lora_adapter_metadata: LoRA adapter metadata associated with the second text encoder to be serialized with the state dict. """ - state_dict = {} - lora_adapter_metadata = {} - - if not (transformer_lora_layers or text_encoder_lora_layers or text_encoder_2_lora_layers): - raise ValueError( - "You must pass at least one of `transformer_lora_layers`, `text_encoder_lora_layers`, `text_encoder_2_lora_layers`." - ) + lora_layers = {} + lora_metadata = {} if transformer_lora_layers: - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) + lora_layers[cls.transformer_name] = transformer_lora_layers + lora_metadata[cls.transformer_name] = transformer_lora_adapter_metadata if text_encoder_lora_layers: - state_dict.update(cls.pack_weights(text_encoder_lora_layers, "text_encoder")) + lora_layers["text_encoder"] = text_encoder_lora_layers + lora_metadata["text_encoder"] = text_encoder_lora_adapter_metadata if text_encoder_2_lora_layers: - state_dict.update(cls.pack_weights(text_encoder_2_lora_layers, "text_encoder_2")) + lora_layers["text_encoder_2"] = text_encoder_2_lora_layers + lora_metadata["text_encoder_2"] = text_encoder_2_lora_adapter_metadata - if transformer_lora_adapter_metadata is not None: - lora_adapter_metadata.update( - _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name) + if not lora_layers: + raise ValueError( + "You must pass at least one of `transformer_lora_layers`, `text_encoder_lora_layers`, or `text_encoder_2_lora_layers`." ) - if text_encoder_lora_adapter_metadata: - lora_adapter_metadata.update( - _pack_dict_with_prefix(text_encoder_lora_adapter_metadata, cls.text_encoder_name) - ) - - if text_encoder_2_lora_adapter_metadata: - lora_adapter_metadata.update( - _pack_dict_with_prefix(text_encoder_2_lora_adapter_metadata, "text_encoder_2") - ) - - cls.write_lora_layers( - state_dict=state_dict, + cls._save_lora_weights( save_directory=save_directory, + lora_layers=lora_layers, + lora_metadata=lora_metadata, is_main_process=is_main_process, weight_name=weight_name, save_function=save_function, safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, ) # Copied from diffusers.loaders.lora_pipeline.StableDiffusionXLLoraLoaderMixin.fuse_lora with unet->transformer @@ -1830,28 +1801,24 @@ class AuraFlowLoraLoaderMixin(LoraBaseMixin): transformer_lora_adapter_metadata: LoRA adapter metadata associated with the transformer to be serialized with the state dict. """ - state_dict = {} - lora_adapter_metadata = {} + lora_layers = {} + lora_metadata = {} - if not transformer_lora_layers: - raise ValueError("You must pass `transformer_lora_layers`.") + if transformer_lora_layers: + lora_layers[cls.transformer_name] = transformer_lora_layers + lora_metadata[cls.transformer_name] = transformer_lora_adapter_metadata - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) + if not lora_layers: + raise ValueError("You must pass at least one of `transformer_lora_layers` or `text_encoder_lora_layers`.") - if transformer_lora_adapter_metadata is not None: - lora_adapter_metadata.update( - _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name) - ) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, + cls._save_lora_weights( save_directory=save_directory, + lora_layers=lora_layers, + lora_metadata=lora_metadata, is_main_process=is_main_process, weight_name=weight_name, save_function=save_function, safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, ) # Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.fuse_lora @@ -2435,37 +2402,28 @@ class FluxLoraLoaderMixin(LoraBaseMixin): text_encoder_lora_adapter_metadata: LoRA adapter metadata associated with the text encoder to be serialized with the state dict. """ - state_dict = {} - lora_adapter_metadata = {} - - if not (transformer_lora_layers or text_encoder_lora_layers): - raise ValueError("You must pass at least one of `transformer_lora_layers` and `text_encoder_lora_layers`.") + lora_layers = {} + lora_metadata = {} if transformer_lora_layers: - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) + lora_layers[cls.transformer_name] = transformer_lora_layers + lora_metadata[cls.transformer_name] = transformer_lora_adapter_metadata if text_encoder_lora_layers: - state_dict.update(cls.pack_weights(text_encoder_lora_layers, cls.text_encoder_name)) + lora_layers[cls.text_encoder_name] = text_encoder_lora_layers + lora_metadata[cls.text_encoder_name] = text_encoder_lora_adapter_metadata - if transformer_lora_adapter_metadata: - lora_adapter_metadata.update( - _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name) - ) + if not lora_layers: + raise ValueError("You must pass at least one of `transformer_lora_layers` or `text_encoder_lora_layers`.") - if text_encoder_lora_adapter_metadata: - lora_adapter_metadata.update( - _pack_dict_with_prefix(text_encoder_lora_adapter_metadata, cls.text_encoder_name) - ) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, + cls._save_lora_weights( save_directory=save_directory, + lora_layers=lora_layers, + lora_metadata=lora_metadata, is_main_process=is_main_process, weight_name=weight_name, save_function=save_function, safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, ) def fuse_lora( @@ -3254,28 +3212,24 @@ class CogVideoXLoraLoaderMixin(LoraBaseMixin): transformer_lora_adapter_metadata: LoRA adapter metadata associated with the transformer to be serialized with the state dict. """ - state_dict = {} - lora_adapter_metadata = {} + lora_layers = {} + lora_metadata = {} - if not transformer_lora_layers: - raise ValueError("You must pass `transformer_lora_layers`.") + if transformer_lora_layers: + lora_layers[cls.transformer_name] = transformer_lora_layers + lora_metadata[cls.transformer_name] = transformer_lora_adapter_metadata - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) + if not lora_layers: + raise ValueError("You must pass at least one of `transformer_lora_layers` or `text_encoder_lora_layers`.") - if transformer_lora_adapter_metadata is not None: - lora_adapter_metadata.update( - _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name) - ) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, + cls._save_lora_weights( save_directory=save_directory, + lora_layers=lora_layers, + lora_metadata=lora_metadata, is_main_process=is_main_process, weight_name=weight_name, save_function=save_function, safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, ) def fuse_lora( @@ -3594,28 +3548,24 @@ class Mochi1LoraLoaderMixin(LoraBaseMixin): transformer_lora_adapter_metadata: LoRA adapter metadata associated with the transformer to be serialized with the state dict. """ - state_dict = {} - lora_adapter_metadata = {} + lora_layers = {} + lora_metadata = {} - if not transformer_lora_layers: - raise ValueError("You must pass `transformer_lora_layers`.") + if transformer_lora_layers: + lora_layers[cls.transformer_name] = transformer_lora_layers + lora_metadata[cls.transformer_name] = transformer_lora_adapter_metadata - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) + if not lora_layers: + raise ValueError("You must pass at least one of `transformer_lora_layers` or `text_encoder_lora_layers`.") - if transformer_lora_adapter_metadata is not None: - lora_adapter_metadata.update( - _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name) - ) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, + cls._save_lora_weights( save_directory=save_directory, + lora_layers=lora_layers, + lora_metadata=lora_metadata, is_main_process=is_main_process, weight_name=weight_name, save_function=save_function, safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, ) # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora @@ -3938,28 +3888,24 @@ class LTXVideoLoraLoaderMixin(LoraBaseMixin): transformer_lora_adapter_metadata: LoRA adapter metadata associated with the transformer to be serialized with the state dict. """ - state_dict = {} - lora_adapter_metadata = {} + lora_layers = {} + lora_metadata = {} - if not transformer_lora_layers: - raise ValueError("You must pass `transformer_lora_layers`.") + if transformer_lora_layers: + lora_layers[cls.transformer_name] = transformer_lora_layers + lora_metadata[cls.transformer_name] = transformer_lora_adapter_metadata - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) + if not lora_layers: + raise ValueError("You must pass at least one of `transformer_lora_layers` or `text_encoder_lora_layers`.") - if transformer_lora_adapter_metadata is not None: - lora_adapter_metadata.update( - _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name) - ) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, + cls._save_lora_weights( save_directory=save_directory, + lora_layers=lora_layers, + lora_metadata=lora_metadata, is_main_process=is_main_process, weight_name=weight_name, save_function=save_function, safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, ) # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora @@ -4280,28 +4226,24 @@ class SanaLoraLoaderMixin(LoraBaseMixin): transformer_lora_adapter_metadata: LoRA adapter metadata associated with the transformer to be serialized with the state dict. """ - state_dict = {} - lora_adapter_metadata = {} + lora_layers = {} + lora_metadata = {} - if not transformer_lora_layers: - raise ValueError("You must pass `transformer_lora_layers`.") + if transformer_lora_layers: + lora_layers[cls.transformer_name] = transformer_lora_layers + lora_metadata[cls.transformer_name] = transformer_lora_adapter_metadata - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) + if not lora_layers: + raise ValueError("You must pass at least one of `transformer_lora_layers` or `text_encoder_lora_layers`.") - if transformer_lora_adapter_metadata is not None: - lora_adapter_metadata.update( - _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name) - ) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, + cls._save_lora_weights( save_directory=save_directory, + lora_layers=lora_layers, + lora_metadata=lora_metadata, is_main_process=is_main_process, weight_name=weight_name, save_function=save_function, safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, ) # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora @@ -4624,28 +4566,24 @@ class HunyuanVideoLoraLoaderMixin(LoraBaseMixin): transformer_lora_adapter_metadata: LoRA adapter metadata associated with the transformer to be serialized with the state dict. """ - state_dict = {} - lora_adapter_metadata = {} + lora_layers = {} + lora_metadata = {} - if not transformer_lora_layers: - raise ValueError("You must pass `transformer_lora_layers`.") + if transformer_lora_layers: + lora_layers[cls.transformer_name] = transformer_lora_layers + lora_metadata[cls.transformer_name] = transformer_lora_adapter_metadata - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) + if not lora_layers: + raise ValueError("You must pass at least one of `transformer_lora_layers` or `text_encoder_lora_layers`.") - if transformer_lora_adapter_metadata is not None: - lora_adapter_metadata.update( - _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name) - ) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, + cls._save_lora_weights( save_directory=save_directory, + lora_layers=lora_layers, + lora_metadata=lora_metadata, is_main_process=is_main_process, weight_name=weight_name, save_function=save_function, safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, ) # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora @@ -4969,28 +4907,24 @@ class Lumina2LoraLoaderMixin(LoraBaseMixin): transformer_lora_adapter_metadata: LoRA adapter metadata associated with the transformer to be serialized with the state dict. """ - state_dict = {} - lora_adapter_metadata = {} + lora_layers = {} + lora_metadata = {} - if not transformer_lora_layers: - raise ValueError("You must pass `transformer_lora_layers`.") + if transformer_lora_layers: + lora_layers[cls.transformer_name] = transformer_lora_layers + lora_metadata[cls.transformer_name] = transformer_lora_adapter_metadata - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) + if not lora_layers: + raise ValueError("You must pass at least one of `transformer_lora_layers` or `text_encoder_lora_layers`.") - if transformer_lora_adapter_metadata is not None: - lora_adapter_metadata.update( - _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name) - ) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, + cls._save_lora_weights( save_directory=save_directory, + lora_layers=lora_layers, + lora_metadata=lora_metadata, is_main_process=is_main_process, weight_name=weight_name, save_function=save_function, safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, ) # Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.fuse_lora @@ -5384,28 +5318,24 @@ class WanLoraLoaderMixin(LoraBaseMixin): transformer_lora_adapter_metadata: LoRA adapter metadata associated with the transformer to be serialized with the state dict. """ - state_dict = {} - lora_adapter_metadata = {} + lora_layers = {} + lora_metadata = {} - if not transformer_lora_layers: - raise ValueError("You must pass `transformer_lora_layers`.") + if transformer_lora_layers: + lora_layers[cls.transformer_name] = transformer_lora_layers + lora_metadata[cls.transformer_name] = transformer_lora_adapter_metadata - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) + if not lora_layers: + raise ValueError("You must pass at least one of `transformer_lora_layers` or `text_encoder_lora_layers`.") - if transformer_lora_adapter_metadata is not None: - lora_adapter_metadata.update( - _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name) - ) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, + cls._save_lora_weights( save_directory=save_directory, + lora_layers=lora_layers, + lora_metadata=lora_metadata, is_main_process=is_main_process, weight_name=weight_name, save_function=save_function, safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, ) # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora @@ -5802,28 +5732,24 @@ class SkyReelsV2LoraLoaderMixin(LoraBaseMixin): transformer_lora_adapter_metadata: LoRA adapter metadata associated with the transformer to be serialized with the state dict. """ - state_dict = {} - lora_adapter_metadata = {} + lora_layers = {} + lora_metadata = {} - if not transformer_lora_layers: - raise ValueError("You must pass `transformer_lora_layers`.") + if transformer_lora_layers: + lora_layers[cls.transformer_name] = transformer_lora_layers + lora_metadata[cls.transformer_name] = transformer_lora_adapter_metadata - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) + if not lora_layers: + raise ValueError("You must pass at least one of `transformer_lora_layers` or `text_encoder_lora_layers`.") - if transformer_lora_adapter_metadata is not None: - lora_adapter_metadata.update( - _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name) - ) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, + cls._save_lora_weights( save_directory=save_directory, + lora_layers=lora_layers, + lora_metadata=lora_metadata, is_main_process=is_main_process, weight_name=weight_name, save_function=save_function, safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, ) # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora @@ -6144,28 +6070,24 @@ class CogView4LoraLoaderMixin(LoraBaseMixin): transformer_lora_adapter_metadata: LoRA adapter metadata associated with the transformer to be serialized with the state dict. """ - state_dict = {} - lora_adapter_metadata = {} + lora_layers = {} + lora_metadata = {} - if not transformer_lora_layers: - raise ValueError("You must pass `transformer_lora_layers`.") + if transformer_lora_layers: + lora_layers[cls.transformer_name] = transformer_lora_layers + lora_metadata[cls.transformer_name] = transformer_lora_adapter_metadata - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) + if not lora_layers: + raise ValueError("You must pass at least one of `transformer_lora_layers` or `text_encoder_lora_layers`.") - if transformer_lora_adapter_metadata is not None: - lora_adapter_metadata.update( - _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name) - ) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, + cls._save_lora_weights( save_directory=save_directory, + lora_layers=lora_layers, + lora_metadata=lora_metadata, is_main_process=is_main_process, weight_name=weight_name, save_function=save_function, safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, ) # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora @@ -6488,28 +6410,24 @@ class HiDreamImageLoraLoaderMixin(LoraBaseMixin): transformer_lora_adapter_metadata: LoRA adapter metadata associated with the transformer to be serialized with the state dict. """ - state_dict = {} - lora_adapter_metadata = {} + lora_layers = {} + lora_metadata = {} - if not transformer_lora_layers: - raise ValueError("You must pass `transformer_lora_layers`.") + if transformer_lora_layers: + lora_layers[cls.transformer_name] = transformer_lora_layers + lora_metadata[cls.transformer_name] = transformer_lora_adapter_metadata - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) + if not lora_layers: + raise ValueError("You must pass at least one of `transformer_lora_layers` or `text_encoder_lora_layers`.") - if transformer_lora_adapter_metadata is not None: - lora_adapter_metadata.update( - _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name) - ) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, + cls._save_lora_weights( save_directory=save_directory, + lora_layers=lora_layers, + lora_metadata=lora_metadata, is_main_process=is_main_process, weight_name=weight_name, save_function=save_function, safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, ) # Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.fuse_lora @@ -6835,28 +6753,24 @@ class QwenImageLoraLoaderMixin(LoraBaseMixin): transformer_lora_adapter_metadata: LoRA adapter metadata associated with the transformer to be serialized with the state dict. """ - state_dict = {} - lora_adapter_metadata = {} + lora_layers = {} + lora_metadata = {} - if not transformer_lora_layers: - raise ValueError("You must pass `transformer_lora_layers`.") + if transformer_lora_layers: + lora_layers[cls.transformer_name] = transformer_lora_layers + lora_metadata[cls.transformer_name] = transformer_lora_adapter_metadata - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) + if not lora_layers: + raise ValueError("You must pass at least one of `transformer_lora_layers` or `text_encoder_lora_layers`.") - if transformer_lora_adapter_metadata is not None: - lora_adapter_metadata.update( - _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name) - ) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, + cls._save_lora_weights( save_directory=save_directory, + lora_layers=lora_layers, + lora_metadata=lora_metadata, is_main_process=is_main_process, weight_name=weight_name, save_function=save_function, safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, ) # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora From 579673501535753e382974e1037498cf61c3c7a1 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Mon, 22 Sep 2025 21:57:30 +0530 Subject: [PATCH 014/123] add test and doc for QwenImageEdit Plus (#12363) * up * xfail some tests * up * up --- docs/source/en/api/pipelines/qwenimage.md | 34 ++- .../qwenimage/test_qwenimage_edit_plus.py | 253 ++++++++++++++++++ 2 files changed, 286 insertions(+), 1 deletion(-) create mode 100644 tests/pipelines/qwenimage/test_qwenimage_edit_plus.py diff --git a/docs/source/en/api/pipelines/qwenimage.md b/docs/source/en/api/pipelines/qwenimage.md index 2dec47309c..4c999bca35 100644 --- a/docs/source/en/api/pipelines/qwenimage.md +++ b/docs/source/en/api/pipelines/qwenimage.md @@ -26,6 +26,7 @@ Qwen-Image comes in the following variants: |:----------:|:--------:| | Qwen-Image | [`Qwen/Qwen-Image`](https://huggingface.co/Qwen/Qwen-Image) | | Qwen-Image-Edit | [`Qwen/Qwen-Image-Edit`](https://huggingface.co/Qwen/Qwen-Image-Edit) | +| Qwen-Image-Edit Plus | [Qwen/Qwen-Image-Edit-2509](https://huggingface.co/Qwen/Qwen-Image-Edit-2509) | @@ -96,6 +97,29 @@ The `guidance_scale` parameter in the pipeline is there to support future guidan +## Multi-image reference with QwenImageEditPlusPipeline + +With [`QwenImageEditPlusPipeline`], one can provide multiple images as input reference. + +``` +import torch +from PIL import Image +from diffusers import QwenImageEditPlusPipeline +from diffusers.utils import load_image + +pipe = QwenImageEditPlusPipeline.from_pretrained( + "Qwen/Qwen-Image-Edit-2509", torch_dtype=torch.bfloat16 +).to("cuda") + +image_1 = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/grumpy.jpg") +image_2 = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/peng.png") +image = pipe( + image=[image_1, image_2], + prompt="put the penguin and the cat at a game show called "Qwen Edit Plus Games"", + num_inference_steps=50 +).images[0] +``` + ## QwenImagePipeline [[autodoc]] QwenImagePipeline @@ -126,7 +150,15 @@ The `guidance_scale` parameter in the pipeline is there to support future guidan - all - __call__ -## QwenImaggeControlNetPipeline +## QwenImageControlNetPipeline + +[[autodoc]] QwenImageControlNetPipeline + - all + - __call__ + +## QwenImageEditPlusPipeline + +[[autodoc]] QwenImageEditPlusPipeline - all - __call__ diff --git a/tests/pipelines/qwenimage/test_qwenimage_edit_plus.py b/tests/pipelines/qwenimage/test_qwenimage_edit_plus.py new file mode 100644 index 0000000000..6faf347282 --- /dev/null +++ b/tests/pipelines/qwenimage/test_qwenimage_edit_plus.py @@ -0,0 +1,253 @@ +# Copyright 2025 The HuggingFace Team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import numpy as np +import pytest +import torch +from PIL import Image +from transformers import Qwen2_5_VLConfig, Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessor + +from diffusers import ( + AutoencoderKLQwenImage, + FlowMatchEulerDiscreteScheduler, + QwenImageEditPlusPipeline, + QwenImageTransformer2DModel, +) + +from ...testing_utils import enable_full_determinism, torch_device +from ..pipeline_params import TEXT_TO_IMAGE_PARAMS +from ..test_pipelines_common import PipelineTesterMixin, to_np + + +enable_full_determinism() + + +class QwenImageEditPlusPipelineFastTests(PipelineTesterMixin, unittest.TestCase): + pipeline_class = QwenImageEditPlusPipeline + params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} + batch_params = frozenset(["prompt", "image"]) + image_params = frozenset(["image"]) + image_latents_params = frozenset(["latents"]) + required_optional_params = frozenset( + [ + "num_inference_steps", + "generator", + "latents", + "return_dict", + "callback_on_step_end", + "callback_on_step_end_tensor_inputs", + ] + ) + supports_dduf = False + test_xformers_attention = False + test_layerwise_casting = True + test_group_offloading = True + + def get_dummy_components(self): + tiny_ckpt_id = "hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration" + + torch.manual_seed(0) + transformer = QwenImageTransformer2DModel( + patch_size=2, + in_channels=16, + out_channels=4, + num_layers=2, + attention_head_dim=16, + num_attention_heads=3, + joint_attention_dim=16, + guidance_embeds=False, + axes_dims_rope=(8, 4, 4), + ) + + torch.manual_seed(0) + z_dim = 4 + vae = AutoencoderKLQwenImage( + base_dim=z_dim * 6, + z_dim=z_dim, + dim_mult=[1, 2, 4], + num_res_blocks=1, + temperal_downsample=[False, True], + latents_mean=[0.0] * z_dim, + latents_std=[1.0] * z_dim, + ) + + torch.manual_seed(0) + scheduler = FlowMatchEulerDiscreteScheduler() + + torch.manual_seed(0) + config = Qwen2_5_VLConfig( + text_config={ + "hidden_size": 16, + "intermediate_size": 16, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 2, + "rope_scaling": { + "mrope_section": [1, 1, 2], + "rope_type": "default", + "type": "default", + }, + "rope_theta": 1000000.0, + }, + vision_config={ + "depth": 2, + "hidden_size": 16, + "intermediate_size": 16, + "num_heads": 2, + "out_hidden_size": 16, + }, + hidden_size=16, + vocab_size=152064, + vision_end_token_id=151653, + vision_start_token_id=151652, + vision_token_id=151654, + ) + text_encoder = Qwen2_5_VLForConditionalGeneration(config) + tokenizer = Qwen2Tokenizer.from_pretrained(tiny_ckpt_id) + + components = { + "transformer": transformer, + "vae": vae, + "scheduler": scheduler, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + "processor": Qwen2VLProcessor.from_pretrained(tiny_ckpt_id), + } + return components + + def get_dummy_inputs(self, device, seed=0): + if str(device).startswith("mps"): + generator = torch.manual_seed(seed) + else: + generator = torch.Generator(device=device).manual_seed(seed) + + image = Image.new("RGB", (32, 32)) + inputs = { + "prompt": "dance monkey", + "image": [image, image], + "negative_prompt": "bad quality", + "generator": generator, + "num_inference_steps": 2, + "true_cfg_scale": 1.0, + "height": 32, + "width": 32, + "max_sequence_length": 16, + "output_type": "pt", + } + + return inputs + + def test_inference(self): + device = "cpu" + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + image = pipe(**inputs).images + generated_image = image[0] + self.assertEqual(generated_image.shape, (3, 32, 32)) + + # fmt: off + expected_slice = torch.tensor([[0.5637, 0.6341, 0.6001, 0.5620, 0.5794, 0.5498, 0.5757, 0.6389, 0.4174, 0.3597, 0.5649, 0.4894, 0.4969, 0.5255, 0.4083, 0.4986]]) + # fmt: on + + generated_slice = generated_image.flatten() + generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) + self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=1e-3)) + + def test_attention_slicing_forward_pass( + self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3 + ): + if not self.test_attention_slicing: + return + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + for component in pipe.components.values(): + if hasattr(component, "set_default_attn_processor"): + component.set_default_attn_processor() + pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + + generator_device = "cpu" + inputs = self.get_dummy_inputs(generator_device) + output_without_slicing = pipe(**inputs)[0] + + pipe.enable_attention_slicing(slice_size=1) + inputs = self.get_dummy_inputs(generator_device) + output_with_slicing1 = pipe(**inputs)[0] + + pipe.enable_attention_slicing(slice_size=2) + inputs = self.get_dummy_inputs(generator_device) + output_with_slicing2 = pipe(**inputs)[0] + + if test_max_difference: + max_diff1 = np.abs(to_np(output_with_slicing1) - to_np(output_without_slicing)).max() + max_diff2 = np.abs(to_np(output_with_slicing2) - to_np(output_without_slicing)).max() + self.assertLess( + max(max_diff1, max_diff2), + expected_max_diff, + "Attention slicing should not affect the inference results", + ) + + def test_vae_tiling(self, expected_diff_max: float = 0.2): + generator_device = "cpu" + components = self.get_dummy_components() + + pipe = self.pipeline_class(**components) + pipe.to("cpu") + pipe.set_progress_bar_config(disable=None) + + # Without tiling + inputs = self.get_dummy_inputs(generator_device) + inputs["height"] = inputs["width"] = 128 + output_without_tiling = pipe(**inputs)[0] + + # With tiling + pipe.vae.enable_tiling( + tile_sample_min_height=96, + tile_sample_min_width=96, + tile_sample_stride_height=64, + tile_sample_stride_width=64, + ) + inputs = self.get_dummy_inputs(generator_device) + inputs["height"] = inputs["width"] = 128 + output_with_tiling = pipe(**inputs)[0] + + self.assertLess( + (to_np(output_without_tiling) - to_np(output_with_tiling)).max(), + expected_diff_max, + "VAE tiling should not affect the inference results", + ) + + @pytest.mark.xfail(condition=True, reason="Preconfigured embeddings need to be revisited.", strict=True) + def test_encode_prompt_works_in_isolation(self, extra_required_param_value_dict=None, atol=1e-4, rtol=1e-4): + super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict, atol, rtol) + + @pytest.mark.xfail(condition=True, reason="Batch of multiple images needs to be revisited", strict=True) + def test_num_images_per_prompt(): + super().test_num_images_per_prompt() + + @pytest.mark.xfail(condition=True, reason="Batch of multiple images needs to be revisited", strict=True) + def test_inference_batch_consistent(): + super().test_inference_batch_consistent() + + @pytest.mark.xfail(condition=True, reason="Batch of multiple images needs to be revisited", strict=True) + def test_inference_batch_single_identical(): + super().test_inference_batch_single_identical() From 1448b035859dd57bbb565239dcdd79a025a85422 Mon Sep 17 00:00:00 2001 From: SahilCarterr <110806554+SahilCarterr@users.noreply.github.com> Date: Tue, 23 Sep 2025 01:34:13 +0530 Subject: [PATCH 015/123] [Fix] chroma docs (#12360) * Fixes chroma docs * fix docs fixed docs are now consistent --- src/diffusers/pipelines/chroma/pipeline_chroma.py | 10 +++++----- .../pipelines/chroma/pipeline_chroma_img2img.py | 12 ++++++------ 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/src/diffusers/pipelines/chroma/pipeline_chroma.py b/src/diffusers/pipelines/chroma/pipeline_chroma.py index f3ed700bc4..19ea7729c9 100644 --- a/src/diffusers/pipelines/chroma/pipeline_chroma.py +++ b/src/diffusers/pipelines/chroma/pipeline_chroma.py @@ -688,11 +688,11 @@ class ChromaPipeline( their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. guidance_scale (`float`, *optional*, defaults to 3.5): - Embedded guiddance scale is enabled by setting `guidance_scale` > 1. Higher `guidance_scale` encourages - a model to generate images more aligned with `prompt` at the expense of lower image quality. - - Guidance-distilled models approximates true classifer-free guidance for `guidance_scale` > 1. Refer to - the [paper](https://huggingface.co/papers/2210.03142) to learn more. + Guidance scale as defined in [Classifier-Free Diffusion + Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. + of [Imagen Paper](https://huggingface.co/papers/2205.11487). 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. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. generator (`torch.Generator` or `List[torch.Generator]`, *optional*): diff --git a/src/diffusers/pipelines/chroma/pipeline_chroma_img2img.py b/src/diffusers/pipelines/chroma/pipeline_chroma_img2img.py index 26f13fe06c..9afd4b9e15 100644 --- a/src/diffusers/pipelines/chroma/pipeline_chroma_img2img.py +++ b/src/diffusers/pipelines/chroma/pipeline_chroma_img2img.py @@ -749,12 +749,12 @@ class ChromaImg2ImgPipeline( Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. - guidance_scale (`float`, *optional*, defaults to 5.0): - Embedded guiddance scale is enabled by setting `guidance_scale` > 1. Higher `guidance_scale` encourages - a model to generate images more aligned with `prompt` at the expense of lower image quality. - - Guidance-distilled models approximates true classifer-free guidance for `guidance_scale` > 1. Refer to - the [paper](https://huggingface.co/papers/2210.03142) to learn more. + guidance_scale (`float`, *optional*, defaults to 3.5): + Guidance scale as defined in [Classifier-Free Diffusion + Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. + of [Imagen Paper](https://huggingface.co/papers/2205.11487). 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. strength (`float, *optional*, defaults to 0.9): Conceptually, indicates how much to transform the reference image. 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 From 76810eca2bbc5ce1e47997e7d9ff690a003fa70f Mon Sep 17 00:00:00 2001 From: Steven Liu <59462357+stevhliu@users.noreply.github.com> Date: Tue, 23 Sep 2025 10:29:16 -0700 Subject: [PATCH 016/123] [docs] Schedulers (#12246) * init * toctree * scheduler suggestions * toctree --- docs/source/en/_toctree.yml | 6 +- docs/source/en/using-diffusers/models.md | 120 ------ .../en/using-diffusers/scheduler_features.md | 235 ----------- docs/source/en/using-diffusers/schedulers.md | 397 +++++++++++------- 4 files changed, 236 insertions(+), 522 deletions(-) delete mode 100644 docs/source/en/using-diffusers/models.md delete mode 100644 docs/source/en/using-diffusers/scheduler_features.md diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 14dbfe3ea1..856874d519 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -23,11 +23,7 @@ - local: using-diffusers/reusing_seeds title: Reproducibility - local: using-diffusers/schedulers - title: Load schedulers and models - - local: using-diffusers/models - title: Models - - local: using-diffusers/scheduler_features - title: Scheduler features + title: Schedulers - local: using-diffusers/other-formats title: Model files and layouts - local: using-diffusers/push_to_hub diff --git a/docs/source/en/using-diffusers/models.md b/docs/source/en/using-diffusers/models.md deleted file mode 100644 index 22c78d490a..0000000000 --- a/docs/source/en/using-diffusers/models.md +++ /dev/null @@ -1,120 +0,0 @@ - - -[[open-in-colab]] - -# Models - -A diffusion model relies on a few individual models working together to generate an output. These models are responsible for denoising, encoding inputs, and decoding latents into the actual outputs. - -This guide will show you how to load models. - -## Loading a model - -All models are loaded with the [`~ModelMixin.from_pretrained`] method, which downloads and caches the latest model version. If the latest files are available in the local cache, [`~ModelMixin.from_pretrained`] reuses files in the cache. - -Pass the `subfolder` argument to [`~ModelMixin.from_pretrained`] to specify where to load the model weights from. Omit the `subfolder` argument if the repository doesn't have a subfolder structure or if you're loading a standalone model. - -```py -from diffusers import QwenImageTransformer2DModel - -model = QwenImageTransformer2DModel.from_pretrained("Qwen/Qwen-Image", subfolder="transformer") -``` - -## AutoModel - -[`AutoModel`] detects the model class from a `model_index.json` file or a model's `config.json` file. It fetches the correct model class from these files and delegates the actual loading to the model class. [`AutoModel`] is useful for automatic model type detection without needing to know the exact model class beforehand. - -```py -from diffusers import AutoModel - -model = AutoModel.from_pretrained( - "Qwen/Qwen-Image", subfolder="transformer" -) -``` - -## Model data types - -Use the `torch_dtype` argument in [`~ModelMixin.from_pretrained`] to load a model with a specific data type. This allows you to load a model in a lower precision to reduce memory usage. - -```py -import torch -from diffusers import QwenImageTransformer2DModel - -model = QwenImageTransformer2DModel.from_pretrained( - "Qwen/Qwen-Image", - subfolder="transformer", - torch_dtype=torch.bfloat16 -) -``` - -[nn.Module.to](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.to) can also convert to a specific data type on the fly. However, it converts *all* weights to the requested data type unlike `torch_dtype` which respects `_keep_in_fp32_modules`. This argument preserves layers in `torch.float32` for numerical stability and best generation quality (see example [_keep_in_fp32_modules](https://github.com/huggingface/diffusers/blob/f864a9a352fa4a220d860bfdd1782e3e5af96382/src/diffusers/models/transformers/transformer_wan.py#L374)) - -```py -from diffusers import QwenImageTransformer2DModel - -model = QwenImageTransformer2DModel.from_pretrained( - "Qwen/Qwen-Image", subfolder="transformer" -) -model = model.to(dtype=torch.float16) -``` - -## Device placement - -Use the `device_map` argument in [`~ModelMixin.from_pretrained`] to place a model on an accelerator like a GPU. It is especially helpful where there are multiple GPUs. - -Diffusers currently provides three options to `device_map` for individual models, `"cuda"`, `"balanced"` and `"auto"`. Refer to the table below to compare the three placement strategies. - -| parameter | description | -|---|---| -| `"cuda"` | places pipeline on a supported accelerator (CUDA) | -| `"balanced"` | evenly distributes pipeline on all GPUs | -| `"auto"` | distribute model from fastest device first to slowest | - -Use the `max_memory` argument in [`~ModelMixin.from_pretrained`] to allocate a maximum amount of memory to use on each device. By default, Diffusers uses the maximum amount available. - -```py -import torch -from diffusers import QwenImagePipeline - -max_memory = {0: "16GB", 1: "16GB"} -pipeline = QwenImagePipeline.from_pretrained( - "Qwen/Qwen-Image", - torch_dtype=torch.bfloat16, - device_map="cuda", - max_memory=max_memory -) -``` - -The `hf_device_map` attribute allows you to access and view the `device_map`. - -```py -print(transformer.hf_device_map) -# {'': device(type='cuda')} -``` - -## Saving models - -Save a model with the [`~ModelMixin.save_pretrained`] method. - -```py -from diffusers import QwenImageTransformer2DModel - -model = QwenImageTransformer2DModel.from_pretrained("Qwen/Qwen-Image", subfolder="transformer") -model.save_pretrained("./local/model") -``` - -For large models, it is helpful to use `max_shard_size` to save a model as multiple shards. A shard can be loaded faster and save memory (refer to the [parallel loading](./loading#parallel-loading) docs for more details), especially if there is more than one GPU. - -```py -model.save_pretrained("./local/model", max_shard_size="5GB") -``` diff --git a/docs/source/en/using-diffusers/scheduler_features.md b/docs/source/en/using-diffusers/scheduler_features.md deleted file mode 100644 index f7977d53d5..0000000000 --- a/docs/source/en/using-diffusers/scheduler_features.md +++ /dev/null @@ -1,235 +0,0 @@ - - -# Scheduler features - -The scheduler is an important component of any diffusion model because it controls the entire denoising (or sampling) process. There are many types of schedulers, some are optimized for speed and some for quality. With Diffusers, you can modify the scheduler configuration to use custom noise schedules, sigmas, and rescale the noise schedule. Changing these parameters can have profound effects on inference quality and speed. - -This guide will demonstrate how to use these features to improve inference quality. - -> [!TIP] -> Diffusers currently only supports the `timesteps` and `sigmas` parameters for a select list of schedulers and pipelines. Feel free to open a [feature request](https://github.com/huggingface/diffusers/issues/new/choose) if you want to extend these parameters to a scheduler and pipeline that does not currently support it! - -## Timestep schedules - -The timestep or noise schedule determines the amount of noise at each sampling step. The scheduler uses this to generate an image with the corresponding amount of noise at each step. The timestep schedule is generated from the scheduler's default configuration, but you can customize the scheduler to use new and optimized sampling schedules that aren't in Diffusers yet. - -For example, [Align Your Steps (AYS)](https://research.nvidia.com/labs/toronto-ai/AlignYourSteps/) is a method for optimizing a sampling schedule to generate a high-quality image in as little as 10 steps. The optimal [10-step schedule](https://github.com/huggingface/diffusers/blob/a7bf77fc284810483f1e60afe34d1d27ad91ce2e/src/diffusers/schedulers/scheduling_utils.py#L51) for Stable Diffusion XL is: - -```py -from diffusers.schedulers import AysSchedules - -sampling_schedule = AysSchedules["StableDiffusionXLTimesteps"] -print(sampling_schedule) -"[999, 845, 730, 587, 443, 310, 193, 116, 53, 13]" -``` - -You can use the AYS sampling schedule in a pipeline by passing it to the `timesteps` parameter. - -```py -pipeline = StableDiffusionXLPipeline.from_pretrained( - "SG161222/RealVisXL_V4.0", - torch_dtype=torch.float16, - variant="fp16", -).to("cuda") -pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config, algorithm_type="sde-dpmsolver++") - -prompt = "A cinematic shot of a cute little rabbit wearing a jacket and doing a thumbs up" -generator = torch.Generator(device="cpu").manual_seed(2487854446) -image = pipeline( - prompt=prompt, - negative_prompt="", - generator=generator, - timesteps=sampling_schedule, -).images[0] -``` - -
-
- -
AYS timestep schedule 10 steps
-
-
- -
Linearly-spaced timestep schedule 10 steps
-
-
- -
Linearly-spaced timestep schedule 25 steps
-
-
- -## Timestep spacing - -The way sample steps are selected in the schedule can affect the quality of the generated image, especially with respect to [rescaling the noise schedule](#rescale-noise-schedule), which can enable a model to generate much brighter or darker images. Diffusers provides three timestep spacing methods: - -- `leading` creates evenly spaced steps -- `linspace` includes the first and last steps and evenly selects the remaining intermediate steps -- `trailing` only includes the last step and evenly selects the remaining intermediate steps starting from the end - -It is recommended to use the `trailing` spacing method because it generates higher quality images with more details when there are fewer sample steps. But the difference in quality is not as obvious for more standard sample step values. - -```py -import torch -from diffusers import StableDiffusionXLPipeline, DPMSolverMultistepScheduler - -pipeline = StableDiffusionXLPipeline.from_pretrained( - "SG161222/RealVisXL_V4.0", - torch_dtype=torch.float16, - variant="fp16", -).to("cuda") -pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config, timestep_spacing="trailing") - -prompt = "A cinematic shot of a cute little black cat sitting on a pumpkin at night" -generator = torch.Generator(device="cpu").manual_seed(2487854446) -image = pipeline( - prompt=prompt, - negative_prompt="", - generator=generator, - num_inference_steps=5, -).images[0] -image -``` - -
-
- -
trailing spacing after 5 steps
-
-
- -
leading spacing after 5 steps
-
-
- -## Sigmas - -The `sigmas` parameter is the amount of noise added at each timestep according to the timestep schedule. Like the `timesteps` parameter, you can customize the `sigmas` parameter to control how much noise is added at each step. When you use a custom `sigmas` value, the `timesteps` are calculated from the custom `sigmas` value and the default scheduler configuration is ignored. - -For example, you can manually pass the [sigmas](https://github.com/huggingface/diffusers/blob/6529ee67ec02fcf58d2fd9242164ea002b351d75/src/diffusers/schedulers/scheduling_utils.py#L55) for something like the 10-step AYS schedule from before to the pipeline. - -```py -import torch - -from diffusers import DiffusionPipeline, EulerDiscreteScheduler - -model_id = "stabilityai/stable-diffusion-xl-base-1.0" -pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", - torch_dtype=torch.float16, - variant="fp16", -).to("cuda") -pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config) - -sigmas = [14.615, 6.315, 3.771, 2.181, 1.342, 0.862, 0.555, 0.380, 0.234, 0.113, 0.0] -prompt = "anthropomorphic capybara wearing a suit and working with a computer" -generator = torch.Generator(device='cuda').manual_seed(123) -image = pipeline( - prompt=prompt, - num_inference_steps=10, - sigmas=sigmas, - generator=generator -).images[0] -``` - -When you take a look at the scheduler's `timesteps` parameter, you'll see that it is the same as the AYS timestep schedule because the `timestep` schedule is calculated from the `sigmas`. - -```py -print(f" timesteps: {pipe.scheduler.timesteps}") -"timesteps: tensor([999., 845., 730., 587., 443., 310., 193., 116., 53., 13.], device='cuda:0')" -``` - -### Karras sigmas - -> [!TIP] -> Refer to the scheduler API [overview](../api/schedulers/overview) for a list of schedulers that support Karras sigmas. -> -> Karras sigmas should not be used for models that weren't trained with them. For example, the base Stable Diffusion XL model shouldn't use Karras sigmas but the [DreamShaperXL](https://hf.co/Lykon/dreamshaper-xl-1-0) model can since they are trained with Karras sigmas. - -Karras scheduler's use the timestep schedule and sigmas from the [Elucidating the Design Space of Diffusion-Based Generative Models](https://hf.co/papers/2206.00364) paper. This scheduler variant applies a smaller amount of noise per step as it approaches the end of the sampling process compared to other schedulers, and can increase the level of details in the generated image. - -Enable Karras sigmas by setting `use_karras_sigmas=True` in the scheduler. - -```py -import torch -from diffusers import StableDiffusionXLPipeline, DPMSolverMultistepScheduler - -pipeline = StableDiffusionXLPipeline.from_pretrained( - "SG161222/RealVisXL_V4.0", - torch_dtype=torch.float16, - variant="fp16", -).to("cuda") -pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config, algorithm_type="sde-dpmsolver++", use_karras_sigmas=True) - -prompt = "A cinematic shot of a cute little rabbit wearing a jacket and doing a thumbs up" -generator = torch.Generator(device="cpu").manual_seed(2487854446) -image = pipeline( - prompt=prompt, - negative_prompt="", - generator=generator, -).images[0] -``` - -
-
- -
Karras sigmas enabled
-
-
- -
Karras sigmas disabled
-
-
- -## Rescale noise schedule - -In the [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://hf.co/papers/2305.08891) paper, the authors discovered that common noise schedules allowed some signal to leak into the last timestep. This signal leakage at inference can cause models to only generate images with medium brightness. By enforcing a zero signal-to-noise ratio (SNR) for the timstep schedule and sampling from the last timestep, the model can be improved to generate very bright or dark images. - -> [!TIP] -> For inference, you need a model that has been trained with *v_prediction*. To train your own model with *v_prediction*, add the following flag to the [train_text_to_image.py](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py) or [train_text_to_image_lora.py](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_lora.py) scripts. -> -> ```bash -> --prediction_type="v_prediction" -> ``` - -For example, load the [ptx0/pseudo-journey-v2](https://hf.co/ptx0/pseudo-journey-v2) checkpoint which was trained with `v_prediction` and the [`DDIMScheduler`]. Configure the following parameters in the [`DDIMScheduler`]: - -* `rescale_betas_zero_snr=True` to rescale the noise schedule to zero SNR -* `timestep_spacing="trailing"` to start sampling from the last timestep - -Set `guidance_rescale` in the pipeline to prevent over-exposure. A lower value increases brightness but some of the details may appear washed out. - -```py -from diffusers import DiffusionPipeline, DDIMScheduler - -pipeline = DiffusionPipeline.from_pretrained("ptx0/pseudo-journey-v2", use_safetensors=True) - -pipeline.scheduler = DDIMScheduler.from_config( - pipeline.scheduler.config, rescale_betas_zero_snr=True, timestep_spacing="trailing" -) -pipeline.to("cuda") -prompt = "cinematic photo of a snowy mountain at night with the northern lights aurora borealis overhead, 35mm photograph, film, professional, 4k, highly detailed" -generator = torch.Generator(device="cpu").manual_seed(23) -image = pipeline(prompt, guidance_rescale=0.7, generator=generator).images[0] -image -``` - -
-
- -
default Stable Diffusion v2-1 image
-
-
- -
image with zero SNR and trailing timestep spacing enabled
-
-
diff --git a/docs/source/en/using-diffusers/schedulers.md b/docs/source/en/using-diffusers/schedulers.md index 6d928f8037..0e236e4e3e 100644 --- a/docs/source/en/using-diffusers/schedulers.md +++ b/docs/source/en/using-diffusers/schedulers.md @@ -10,200 +10,273 @@ an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express o specific language governing permissions and limitations under the License. --> -# Load schedulers and models - [[open-in-colab]] -Diffusion pipelines are a collection of interchangeable schedulers and models that can be mixed and matched to tailor a pipeline to a specific use case. The scheduler encapsulates the entire denoising process such as the number of denoising steps and the algorithm for finding the denoised sample. A scheduler is not parameterized or trained so they don't take very much memory. The model is usually only concerned with the forward pass of going from a noisy input to a less noisy sample. +# Schedulers -This guide will show you how to load schedulers and models to customize a pipeline. You'll use the [stable-diffusion-v1-5/stable-diffusion-v1-5](https://hf.co/stable-diffusion-v1-5/stable-diffusion-v1-5) checkpoint throughout this guide, so let's load it first. +A scheduler is an algorithm that provides instructions to the denoising process such as how much noise to remove at a certain step. It takes the model prediction from step *t* and applies an update for how to compute the next sample at step *t-1*. Different schedulers produce different results; some are faster while others are more accurate. + +Diffusers supports many schedulers and allows you to modify their timestep schedules, timestep spacing, and more, to generate high-quality images in fewer steps. + +This guide will show you how to load and customize schedulers. + +## Loading schedulers + +Schedulers don't have any parameters and are defined in a configuration file. Access the `.scheduler` attribute of a pipeline to view the configuration. ```py import torch from diffusers import DiffusionPipeline pipeline = DiffusionPipeline.from_pretrained( - "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True -).to("cuda") -``` - -You can see what scheduler this pipeline uses with the `pipeline.scheduler` attribute. - -```py + "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, device_map="cuda" +) pipeline.scheduler -PNDMScheduler { - "_class_name": "PNDMScheduler", - "_diffusers_version": "0.21.4", - "beta_end": 0.012, - "beta_schedule": "scaled_linear", - "beta_start": 0.00085, - "clip_sample": false, - "num_train_timesteps": 1000, - "set_alpha_to_one": false, - "skip_prk_steps": true, - "steps_offset": 1, - "timestep_spacing": "leading", - "trained_betas": null -} ``` -## Load a scheduler - -Schedulers are defined by a configuration file that can be used by a variety of schedulers. Load a scheduler with the [`SchedulerMixin.from_pretrained`] method, and specify the `subfolder` parameter to load the configuration file into the correct subfolder of the pipeline repository. - -For example, to load the [`DDIMScheduler`]: - -```py -from diffusers import DDIMScheduler, DiffusionPipeline - -ddim = DDIMScheduler.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="scheduler") -``` - -Then you can pass the newly loaded scheduler to the pipeline. - -```python -pipeline = DiffusionPipeline.from_pretrained( - "stable-diffusion-v1-5/stable-diffusion-v1-5", scheduler=ddim, torch_dtype=torch.float16, use_safetensors=True -).to("cuda") -``` - -## Compare schedulers - -Schedulers have their own unique strengths and weaknesses, making it difficult to quantitatively compare which scheduler works best for a pipeline. You typically have to make a trade-off between denoising speed and denoising quality. We recommend trying out different schedulers to find one that works best for your use case. Call the `pipeline.scheduler.compatibles` attribute to see what schedulers are compatible with a pipeline. - -Let's compare the [`LMSDiscreteScheduler`], [`EulerDiscreteScheduler`], [`EulerAncestralDiscreteScheduler`], and the [`DPMSolverMultistepScheduler`] on the following prompt and seed. - -```py -import torch -from diffusers import DiffusionPipeline - -pipeline = DiffusionPipeline.from_pretrained( - "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True -).to("cuda") - -prompt = "A photograph of an astronaut riding a horse on Mars, high resolution, high definition." -generator = torch.Generator(device="cuda").manual_seed(8) -``` - -To change the pipelines scheduler, use the [`~ConfigMixin.from_config`] method to load a different scheduler's `pipeline.scheduler.config` into the pipeline. - - - - -[`LMSDiscreteScheduler`] typically generates higher quality images than the default scheduler. - -```py -from diffusers import LMSDiscreteScheduler - -pipeline.scheduler = LMSDiscreteScheduler.from_config(pipeline.scheduler.config) -image = pipeline(prompt, generator=generator).images[0] -image -``` - - - - -[`EulerDiscreteScheduler`] can generate higher quality images in just 30 steps. - -```py -from diffusers import EulerDiscreteScheduler - -pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config) -image = pipeline(prompt, generator=generator).images[0] -image -``` - - - - -[`EulerAncestralDiscreteScheduler`] can generate higher quality images in just 30 steps. - -```py -from diffusers import EulerAncestralDiscreteScheduler - -pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(pipeline.scheduler.config) -image = pipeline(prompt, generator=generator).images[0] -image -``` - - - - -[`DPMSolverMultistepScheduler`] provides a balance between speed and quality and can generate higher quality images in just 20 steps. +Load a different scheduler with [`~SchedulerMixin.from_pretrained`] and specify the `subfolder` argument to load the configuration file into the correct subfolder of the pipeline repository. Pass the new scheduler to the existing pipeline. ```py from diffusers import DPMSolverMultistepScheduler -pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config) -image = pipeline(prompt, generator=generator).images[0] +dpm = DPMSolverMultistepScheduler.from_pretrained( + "stabilityai/stable-diffusion-xl-base-1.0", subfolder="scheduler" +) +pipeline = DiffusionPipeline.from_pretrained( + "stabilityai/stable-diffusion-xl-base-1.0", + scheduler=dpm, + torch_dtype=torch.float16, + device_map="cuda" +) +pipeline.scheduler +``` + +## Timestep schedules + +Timestep or noise schedule decides how noise is distributed over the denoising process. The schedule can be linear or more concentrated toward the beginning or end. It is a precomputed sequence of noise levels generated from the scheduler's default configuration, but it can be customized to use other schedules. + +> [!TIP] +> The `timesteps` argument is only supported for a select list of schedulers and pipelines. Feel free to open a feature request if you want to extend these parameters to a scheduler and pipeline that does not currently support it! + +The example below uses the [Align Your Steps (AYS)](https://research.nvidia.com/labs/toronto-ai/AlignYourSteps/) schedule which can generate a high-quality image in 10 steps, significantly speeding up generation and reducing computation time. + +Import the schedule and pass it to the `timesteps` argument in the pipeline. + +```py +import torch +from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler +from diffusers.schedulers import AysSchedules + +sampling_schedule = AysSchedules["StableDiffusionXLTimesteps"] +print(sampling_schedule) +"[999, 845, 730, 587, 443, 310, 193, 116, 53, 13]" + +pipeline = DiffusionPipeline.from_pretrained( + "SG161222/RealVisXL_V4.0", + torch_dtype=torch.float16, + device_map="cuda" +) +pipeline.scheduler = DPMSolverMultistepScheduler.from_config( + pipeline.scheduler.config, algorithm_type="sde-dpmsolver++" +) + +prompt = "A cinematic shot of a cute little rabbit wearing a jacket and doing a thumbs up" +image = pipeline( + prompt=prompt, + negative_prompt="", + timesteps=sampling_schedule, +).images[0] +``` + +
+
+ +
AYS timestep schedule 10 steps
+
+
+ +
Linearly-spaced timestep schedule 10 steps
+
+
+ +
Linearly-spaced timestep schedule 25 steps
+
+
+ +### Rescaling schedules + +Denoising should begin with pure noise and the signal-to-noise (SNR) ration should be zero. However, some models don't actually start from pure noise which makes it difficult to generate images at brightness extremes. + +> [!TIP] +> Train your own model with `v_prediction` by adding the `--prediction_type="v_prediction"` flag to your training script. You can also [search](https://huggingface.co/search/full-text?q=v_prediction&type=model) for existing models trained with `v_prediction`. + +To fix this, a model must be trained with `v_prediction`. If a model is trained with `v_prediction`, then enable the following arguments in the scheduler. + +- Set `rescale_betas_zero_snr=True` to rescale the noise schedule to the very last timestep with exactly zero SNR +- Set `timestep_spacing="trailing"` to force sampling from the last timestep with pure noise + +```py +from diffusers import DiffusionPipeline, DDIMScheduler + +pipeline = DiffusionPipeline.from_pretrained("ptx0/pseudo-journey-v2", device_map="cuda") + +pipeline.scheduler = DDIMScheduler.from_config( + pipeline.scheduler.config, rescale_betas_zero_snr=True, timestep_spacing="trailing" +) +``` + +Set `guidance_rescale` in the pipeline to avoid overexposed images. A lower value increases brightness, but some details may appear washed out. + +```py +prompt = """ +cinematic photo of a snowy mountain at night with the northern lights aurora borealis +overhead, 35mm photograph, film, professional, 4k, highly detailed +""" +image = pipeline(prompt, guidance_rescale=0.7).images[0] +``` + +
+
+ +
default Stable Diffusion v2-1 image
+
+
+ +
image with zero SNR and trailing timestep spacing enabled
+
+
+ +## Timestep spacing + +Timestep spacing refers to the specific steps *t* to sample from from the schedule. Diffusers provides three spacing types as shown below. + +| spacing strategy | spacing calculation | example timesteps | +|---|---|---| +| `leading` | evenly spaced steps | `[900, 800, 700, ..., 100, 0]` | +| `linspace` | include first and last steps and evenly divide remaining intermediate steps | `[1000, 888.89, 777.78, ..., 111.11, 0]` | +| `trailing` | include last step and evenly divide remaining intermediate steps beginning from the end | `[999, 899, 799, 699, 599, 499, 399, 299, 199, 99]` | + +Pass the spacing strategy to the `timestep_spacing` argument in the scheduler. + +> [!TIP] +> The `trailing` strategy typically produces higher quality images with more details with fewer steps, but the difference in quality is not as obvious for more standard step values. + +```py +import torch +from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler + +pipeline = DiffusionPipeline.from_pretrained( + "SG161222/RealVisXL_V4.0", + torch_dtype=torch.float16, + device_map="cuda" +) +pipeline.scheduler = DPMSolverMultistepScheduler.from_config( + pipeline.scheduler.config, timestep_spacing="trailing" +) + +prompt = "A cinematic shot of a cute little black cat sitting on a pumpkin at night" +image = pipeline( + prompt=prompt, + negative_prompt="", + num_inference_steps=5, +).images[0] image ``` - - -
- -
LMSDiscreteScheduler
+ +
trailing spacing after 5 steps
- -
EulerDiscreteScheduler
-
-
-
-
- -
EulerAncestralDiscreteScheduler
-
-
- -
DPMSolverMultistepScheduler
+ +
leading spacing after 5 steps
-Most images look very similar and are comparable in quality. Again, it often comes down to your specific use case so a good approach is to run multiple different schedulers and compare the results. +## Sigmas -## Models +Sigmas is a measure of how noisy a sample is at a certain step as defined by the schedule. When using custom `sigmas`, the `timesteps` are calculated from these values instead of the default scheduler configuration. -Models are loaded from the [`ModelMixin.from_pretrained`] method, which downloads and caches the latest version of the model weights and configurations. If the latest files are available in the local cache, [`~ModelMixin.from_pretrained`] reuses files in the cache instead of re-downloading them. +> [!TIP] +> The `sigmas` argument is only supported for a select list of schedulers and pipelines. Feel free to open a feature request if you want to extend these parameters to a scheduler and pipeline that does not currently support it! -Models can be loaded from a subfolder with the `subfolder` argument. For example, the model weights for [stable-diffusion-v1-5/stable-diffusion-v1-5](https://hf.co/stable-diffusion-v1-5/stable-diffusion-v1-5) are stored in the [unet](https://hf.co/stable-diffusion-v1-5/stable-diffusion-v1-5/tree/main/unet) subfolder. - -```python -from diffusers import UNet2DConditionModel - -unet = UNet2DConditionModel.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="unet", use_safetensors=True) -``` - -They can also be directly loaded from a [repository](https://huggingface.co/google/ddpm-cifar10-32/tree/main). - -```python -from diffusers import UNet2DModel - -unet = UNet2DModel.from_pretrained("google/ddpm-cifar10-32", use_safetensors=True) -``` - -To load and save model variants, specify the `variant` argument in [`ModelMixin.from_pretrained`] and [`ModelMixin.save_pretrained`]. - -```python -from diffusers import UNet2DConditionModel - -unet = UNet2DConditionModel.from_pretrained( - "stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="unet", variant="non_ema", use_safetensors=True -) -unet.save_pretrained("./local-unet", variant="non_ema") -``` - -Use the `torch_dtype` argument in [`~ModelMixin.from_pretrained`] to specify the dtype to load a model in. +Pass the custom sigmas to the `sigmas` argument in the pipeline. The example below uses the [sigmas](https://github.com/huggingface/diffusers/blob/6529ee67ec02fcf58d2fd9242164ea002b351d75/src/diffusers/schedulers/scheduling_utils.py#L55) from the 10-step AYS schedule. ```py -from diffusers import AutoModel +import torch +from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler -unet = AutoModel.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet", torch_dtype=torch.float16 +pipeline = DiffusionPipeline.from_pretrained( + "SG161222/RealVisXL_V4.0", + torch_dtype=torch.float16, + device_map="cuda" ) +pipeline.scheduler = DPMSolverMultistepScheduler.from_config( + pipeline.scheduler.config, algorithm_type="sde-dpmsolver++" +) + +sigmas = [14.615, 6.315, 3.771, 2.181, 1.342, 0.862, 0.555, 0.380, 0.234, 0.113, 0.0] +prompt = "A cinematic shot of a cute little rabbit wearing a jacket and doing a thumbs up" +image = pipeline( + prompt=prompt, + negative_prompt="", + sigmas=sigmas, +).images[0] ``` -You can also use the [torch.Tensor.to](https://docs.pytorch.org/docs/stable/generated/torch.Tensor.to.html) method to convert to the specified dtype on the fly. It converts *all* weights unlike the `torch_dtype` argument that respects the `_keep_in_fp32_modules`. This is important for models whose layers must remain in fp32 for numerical stability and best generation quality (see example [here](https://github.com/huggingface/diffusers/blob/f864a9a352fa4a220d860bfdd1782e3e5af96382/src/diffusers/models/transformers/transformer_wan.py#L374)). +### Karras sigmas + +[Karras sigmas](https://huggingface.co/papers/2206.00364) resamples the noise schedule for more efficient sampling by clustering sigmas more densely in the middle of the sequence where structure reconstruction is critical, while using fewer sigmas at the beginning and end where noise changes have less impact. This can increase the level of details in a generated image. + +Set `use_karras_sigmas=True` in the scheduler to enable it. + +```py +import torch +from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler + +pipeline = DiffusionPipeline.from_pretrained( + "SG161222/RealVisXL_V4.0", + torch_dtype=torch.float16, + device_map="cuda" +) +pipeline.scheduler = DPMSolverMultistepScheduler.from_config( + pipeline.scheduler.config, + algorithm_type="sde-dpmsolver++", + use_karras_sigmas=True, +) + +prompt = "A cinematic shot of a cute little rabbit wearing a jacket and doing a thumbs up" +image = pipeline( + prompt=prompt, + negative_prompt="", + sigmas=sigmas, +).images[0] +``` + +
+
+ +
Karras sigmas enabled
+
+
+ +
Karras sigmas disabled
+
+
+ +Refer to the scheduler API [overview](../api/schedulers/overview) for a list of schedulers that support Karras sigmas. It should only be used for models trained with Karras sigmas. + +## Choosing a scheduler + +It's important to try different schedulers to find the best one for your use case. Here are a few recommendations to help you get started. + +- DPM++ 2M SDE Karras is generally a good all-purpose option. +- [`TCDScheduler`] works well for distilled models. +- [`FlowMatchEulerDiscreteScheduler`] and [`FlowMatchHeunDiscreteScheduler`] for FlowMatch models. +- [`EulerDiscreteScheduler`] or [`EulerAncestralDiscreteScheduler`] for generating anime style images. +- DPM++ 2M paired with [`LCMScheduler`] on SDXL for generating realistic images. + +## Resources + +- Read the [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) paper for more details about rescaling the noise schedule to enforce zero SNR. \ No newline at end of file From 80de641c1c7973dd83cdb9ebb0946affb28e00f1 Mon Sep 17 00:00:00 2001 From: Dhruv Nair Date: Tue, 23 Sep 2025 19:31:42 +0200 Subject: [PATCH 017/123] Allow Automodel to support custom model code (#12353) * update * update --- src/diffusers/models/auto_model.py | 67 +++++++++++++------- src/diffusers/utils/dynamic_modules_utils.py | 5 ++ 2 files changed, 50 insertions(+), 22 deletions(-) diff --git a/src/diffusers/models/auto_model.py b/src/diffusers/models/auto_model.py index bfe386f1f6..ada0d54e54 100644 --- a/src/diffusers/models/auto_model.py +++ b/src/diffusers/models/auto_model.py @@ -19,6 +19,7 @@ from huggingface_hub.utils import validate_hf_hub_args from ..configuration_utils import ConfigMixin from ..utils import logging +from ..utils.dynamic_modules_utils import get_class_from_dynamic_module, resolve_trust_remote_code logger = logging.get_logger(__name__) @@ -114,6 +115,8 @@ class AutoModel(ConfigMixin): disable_mmap ('bool', *optional*, defaults to 'False'): Whether to disable mmap when loading a Safetensors model. This option can perform better when the model is on a network mount or hard drive, which may not handle the seeky-ness of mmap very well. + trust_remote_cocde (`bool`, *optional*, defaults to `False`): + Whether to trust remote code @@ -140,22 +143,22 @@ class AutoModel(ConfigMixin): You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference. ``` """ - cache_dir = kwargs.pop("cache_dir", None) - force_download = kwargs.pop("force_download", False) - proxies = kwargs.pop("proxies", None) - token = kwargs.pop("token", None) - local_files_only = kwargs.pop("local_files_only", False) - revision = kwargs.pop("revision", None) subfolder = kwargs.pop("subfolder", None) + trust_remote_code = kwargs.pop("trust_remote_code", False) - load_config_kwargs = { - "cache_dir": cache_dir, - "force_download": force_download, - "proxies": proxies, - "token": token, - "local_files_only": local_files_only, - "revision": revision, - } + hub_kwargs_names = [ + "cache_dir", + "force_download", + "local_files_only", + "proxies", + "resume_download", + "revision", + "token", + ] + hub_kwargs = {name: kwargs.pop(name, None) for name in hub_kwargs_names} + + # load_config_kwargs uses the same hub kwargs minus subfolder and resume_download + load_config_kwargs = {k: v for k, v in hub_kwargs.items() if k not in ["subfolder", "resume_download"]} library = None orig_class_name = None @@ -189,15 +192,35 @@ class AutoModel(ConfigMixin): else: raise ValueError(f"Couldn't find model associated with the config file at {pretrained_model_or_path}.") - from ..pipelines.pipeline_loading_utils import ALL_IMPORTABLE_CLASSES, get_class_obj_and_candidates + has_remote_code = "auto_map" in config and cls.__name__ in config["auto_map"] + trust_remote_code = resolve_trust_remote_code(trust_remote_code, pretrained_model_or_path, has_remote_code) + if not (has_remote_code and trust_remote_code): + raise ValueError( + "Selected model repository does not happear to have any custom code or does not have a valid `config.json` file." + ) - model_cls, _ = get_class_obj_and_candidates( - library_name=library, - class_name=orig_class_name, - importable_classes=ALL_IMPORTABLE_CLASSES, - pipelines=None, - is_pipeline_module=False, - ) + if has_remote_code and trust_remote_code: + class_ref = config["auto_map"][cls.__name__] + module_file, class_name = class_ref.split(".") + module_file = module_file + ".py" + model_cls = get_class_from_dynamic_module( + pretrained_model_or_path, + subfolder=subfolder, + module_file=module_file, + class_name=class_name, + **hub_kwargs, + **kwargs, + ) + else: + from ..pipelines.pipeline_loading_utils import ALL_IMPORTABLE_CLASSES, get_class_obj_and_candidates + + model_cls, _ = get_class_obj_and_candidates( + library_name=library, + class_name=orig_class_name, + importable_classes=ALL_IMPORTABLE_CLASSES, + pipelines=None, + is_pipeline_module=False, + ) if model_cls is None: raise ValueError(f"AutoModel can't find a model linked to {orig_class_name}.") diff --git a/src/diffusers/utils/dynamic_modules_utils.py b/src/diffusers/utils/dynamic_modules_utils.py index 674eb65773..de947a12e2 100644 --- a/src/diffusers/utils/dynamic_modules_utils.py +++ b/src/diffusers/utils/dynamic_modules_utils.py @@ -247,6 +247,7 @@ def find_pipeline_class(loaded_module): def get_cached_module_file( pretrained_model_name_or_path: Union[str, os.PathLike], module_file: str, + subfolder: Optional[str] = None, cache_dir: Optional[Union[str, os.PathLike]] = None, force_download: bool = False, proxies: Optional[Dict[str, str]] = None, @@ -353,6 +354,7 @@ def get_cached_module_file( resolved_module_file = hf_hub_download( pretrained_model_name_or_path, module_file, + subfolder=subfolder, cache_dir=cache_dir, force_download=force_download, proxies=proxies, @@ -410,6 +412,7 @@ def get_cached_module_file( get_cached_module_file( pretrained_model_name_or_path, f"{module_needed}.py", + subfolder=subfolder, cache_dir=cache_dir, force_download=force_download, proxies=proxies, @@ -424,6 +427,7 @@ def get_cached_module_file( def get_class_from_dynamic_module( pretrained_model_name_or_path: Union[str, os.PathLike], module_file: str, + subfolder: Optional[str] = None, class_name: Optional[str] = None, cache_dir: Optional[Union[str, os.PathLike]] = None, force_download: bool = False, @@ -497,6 +501,7 @@ def get_class_from_dynamic_module( final_module = get_cached_module_file( pretrained_model_name_or_path, module_file, + subfolder=subfolder, cache_dir=cache_dir, force_download=force_download, proxies=proxies, From a72bc0c4bb817d382a59b38bba8d9a71661f56cb Mon Sep 17 00:00:00 2001 From: Steven Liu <59462357+stevhliu@users.noreply.github.com> Date: Tue, 23 Sep 2025 10:59:46 -0700 Subject: [PATCH 018/123] [docs] Attention backends (#12320) * init * feedback * update * feedback * fixes --- docs/source/en/_toctree.yml | 2 + .../en/optimization/attention_backends.md | 106 ++++++++++++++++++ 2 files changed, 108 insertions(+) create mode 100644 docs/source/en/optimization/attention_backends.md diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 856874d519..4879a7bf04 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -64,6 +64,8 @@ title: Accelerate inference - local: optimization/cache title: Caching + - local: optimization/attention_backends + title: Attention backends - local: optimization/memory title: Reduce memory usage - local: optimization/speed-memory-optims diff --git a/docs/source/en/optimization/attention_backends.md b/docs/source/en/optimization/attention_backends.md new file mode 100644 index 0000000000..04c8b4ba92 --- /dev/null +++ b/docs/source/en/optimization/attention_backends.md @@ -0,0 +1,106 @@ + + +# Attention backends + +> [!TIP] +> The attention dispatcher is an experimental feature. Please open an issue if you have any feedback or encounter any problems. + +Diffusers provides several optimized attention algorithms that are more memory and computationally efficient through it's *attention dispatcher*. The dispatcher acts as a router for managing and switching between different attention implementations and provides a unified interface for interacting with them. + +Refer to the table below for an overview of the available attention families and to the [Available backends](#available-backends) section for a more complete list. + +| attention family | main feature | +|---|---| +| FlashAttention | minimizes memory reads/writes through tiling and recomputation | +| SageAttention | quantizes attention to int8 | +| PyTorch native | built-in PyTorch implementation using [scaled_dot_product_attention](./fp16#scaled-dot-product-attention) | +| xFormers | memory-efficient attention with support for various attention kernels | + +This guide will show you how to set and use the different attention backends. + +## set_attention_backend + +The [`~ModelMixin.set_attention_backend`] method iterates through all the modules in the model and sets the appropriate attention backend to use. The attention backend setting persists until [`~ModelMixin.reset_attention_backend`] is called. + +The example below demonstrates how to enable the `_flash_3_hub` implementation for FlashAttention-3 from the [kernel](https://github.com/huggingface/kernels) library, which allows you to instantly use optimized compute kernels from the Hub without requiring any setup. + +> [!TIP] +> FlashAttention-3 is not supported for non-Hopper architectures, in which case, use FlashAttention with `set_attention_backend("flash")`. + +```py +import torch +from diffusers import QwenImagePipeline + +pipeline = QwenImagePipeline.from_pretrained( + "Qwen/Qwen-Image", torch_dtype=torch.bfloat16, device_map="cuda" +) +pipeline.transformer.set_attention_backend("_flash_3_hub") + +prompt = """ +cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California +highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain +""" +pipeline(prompt).images[0] +``` + +To restore the default attention backend, call [`~ModelMixin.reset_attention_backend`]. + +```py +pipeline.transformer.reset_attention_backend() +``` + +## attention_backend context manager + +The [attention_backend](https://github.com/huggingface/diffusers/blob/5e181eddfe7e44c1444a2511b0d8e21d177850a0/src/diffusers/models/attention_dispatch.py#L225) context manager temporarily sets an attention backend for a model within the context. Outside the context, the default attention (PyTorch's native scaled dot product attention) is used. This is useful if you want to use different backends for different parts of a pipeline or if you want to test the different backends. + +```py +import torch +from diffusers import QwenImagePipeline + +pipeline = QwenImagePipeline.from_pretrained( + "Qwen/Qwen-Image", torch_dtype=torch.bfloat16, device_map="cuda" +) +prompt = """ +cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California +highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain +""" + +with attention_backend("_flash_3_hub"): + image = pipeline(prompt).images[0] +``` + +## Available backends + +Refer to the table below for a complete list of available attention backends and their variants. + +| Backend Name | Family | Description | +|--------------|--------|-------------| +| `native` | [PyTorch native](https://docs.pytorch.org/docs/stable/generated/torch.nn.attention.SDPBackend.html#torch.nn.attention.SDPBackend) | Default backend using PyTorch's scaled_dot_product_attention | +| `flex` | [FlexAttention](https://docs.pytorch.org/docs/stable/nn.attention.flex_attention.html#module-torch.nn.attention.flex_attention) | PyTorch FlexAttention implementation | +| `_native_cudnn` | [PyTorch native](https://docs.pytorch.org/docs/stable/generated/torch.nn.attention.SDPBackend.html#torch.nn.attention.SDPBackend) | CuDNN-optimized attention | +| `_native_efficient` | [PyTorch native](https://docs.pytorch.org/docs/stable/generated/torch.nn.attention.SDPBackend.html#torch.nn.attention.SDPBackend) | Memory-efficient attention | +| `_native_flash` | [PyTorch native](https://docs.pytorch.org/docs/stable/generated/torch.nn.attention.SDPBackend.html#torch.nn.attention.SDPBackend) | PyTorch's FlashAttention | +| `_native_math` | [PyTorch native](https://docs.pytorch.org/docs/stable/generated/torch.nn.attention.SDPBackend.html#torch.nn.attention.SDPBackend) | Math-based attention (fallback) | +| `_native_npu` | [PyTorch native](https://docs.pytorch.org/docs/stable/generated/torch.nn.attention.SDPBackend.html#torch.nn.attention.SDPBackend) | NPU-optimized attention | +| `_native_xla` | [PyTorch native](https://docs.pytorch.org/docs/stable/generated/torch.nn.attention.SDPBackend.html#torch.nn.attention.SDPBackend) | XLA-optimized attention | +| `flash` | [FlashAttention](https://github.com/Dao-AILab/flash-attention) | FlashAttention-2 | +| `flash_varlen` | [FlashAttention](https://github.com/Dao-AILab/flash-attention) | Variable length FlashAttention | +| `_flash_3` | [FlashAttention](https://github.com/Dao-AILab/flash-attention) | FlashAttention-3 | +| `_flash_varlen_3` | [FlashAttention](https://github.com/Dao-AILab/flash-attention) | Variable length FlashAttention-3 | +| `_flash_3_hub` | [FlashAttention](https://github.com/Dao-AILab/flash-attention) | FlashAttention-3 from kernels | +| `sage` | [SageAttention](https://github.com/thu-ml/SageAttention) | Quantized attention (INT8 QK) | +| `sage_varlen` | [SageAttention](https://github.com/thu-ml/SageAttention) | Variable length SageAttention | +| `_sage_qk_int8_pv_fp8_cuda` | [SageAttention](https://github.com/thu-ml/SageAttention) | INT8 QK + FP8 PV (CUDA) | +| `_sage_qk_int8_pv_fp8_cuda_sm90` | [SageAttention](https://github.com/thu-ml/SageAttention) | INT8 QK + FP8 PV (SM90) | +| `_sage_qk_int8_pv_fp16_cuda` | [SageAttention](https://github.com/thu-ml/SageAttention) | INT8 QK + FP16 PV (CUDA) | +| `_sage_qk_int8_pv_fp16_triton` | [SageAttention](https://github.com/thu-ml/SageAttention) | INT8 QK + FP16 PV (Triton) | +| `xformers` | [xFormers](https://github.com/facebookresearch/xformers) | Memory-efficient attention | From 09e777a3e13cf811e35da57abfe6ce239d9b0f15 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Wed, 24 Sep 2025 08:36:50 +0530 Subject: [PATCH 019/123] [tests] Single scheduler in lora tests (#12315) * single scheduler please. * up * up * up --- tests/lora/test_lora_layers_auraflow.py | 1 - tests/lora/test_lora_layers_cogvideox.py | 2 - tests/lora/test_lora_layers_cogview4.py | 36 +- tests/lora/test_lora_layers_flux.py | 6 +- tests/lora/test_lora_layers_hunyuanvideo.py | 1 - tests/lora/test_lora_layers_ltx_video.py | 1 - tests/lora/test_lora_layers_lumina2.py | 46 +- tests/lora/test_lora_layers_mochi.py | 1 - tests/lora/test_lora_layers_qwenimage.py | 1 - tests/lora/test_lora_layers_sana.py | 5 +- tests/lora/test_lora_layers_sd3.py | 1 - tests/lora/test_lora_layers_wan.py | 1 - tests/lora/test_lora_layers_wanvace.py | 4 +- tests/lora/utils.py | 2083 +++++++++---------- 14 files changed, 1045 insertions(+), 1144 deletions(-) diff --git a/tests/lora/test_lora_layers_auraflow.py b/tests/lora/test_lora_layers_auraflow.py index 67084dd6d0..91f63c4b56 100644 --- a/tests/lora/test_lora_layers_auraflow.py +++ b/tests/lora/test_lora_layers_auraflow.py @@ -43,7 +43,6 @@ from .utils import PeftLoraLoaderMixinTests # noqa: E402 class AuraFlowLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): pipeline_class = AuraFlowPipeline scheduler_cls = FlowMatchEulerDiscreteScheduler - scheduler_classes = [FlowMatchEulerDiscreteScheduler] scheduler_kwargs = {} transformer_kwargs = { diff --git a/tests/lora/test_lora_layers_cogvideox.py b/tests/lora/test_lora_layers_cogvideox.py index 16147f35c7..fa57b4c9c2 100644 --- a/tests/lora/test_lora_layers_cogvideox.py +++ b/tests/lora/test_lora_layers_cogvideox.py @@ -21,7 +21,6 @@ from transformers import AutoTokenizer, T5EncoderModel from diffusers import ( AutoencoderKLCogVideoX, - CogVideoXDDIMScheduler, CogVideoXDPMScheduler, CogVideoXPipeline, CogVideoXTransformer3DModel, @@ -44,7 +43,6 @@ class CogVideoXLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): pipeline_class = CogVideoXPipeline scheduler_cls = CogVideoXDPMScheduler scheduler_kwargs = {"timestep_spacing": "trailing"} - scheduler_classes = [CogVideoXDDIMScheduler, CogVideoXDPMScheduler] transformer_kwargs = { "num_attention_heads": 4, diff --git a/tests/lora/test_lora_layers_cogview4.py b/tests/lora/test_lora_layers_cogview4.py index 3b8a56c403..9c62d2f0b8 100644 --- a/tests/lora/test_lora_layers_cogview4.py +++ b/tests/lora/test_lora_layers_cogview4.py @@ -50,7 +50,6 @@ class TokenizerWrapper: class CogView4LoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): pipeline_class = CogView4Pipeline scheduler_cls = FlowMatchEulerDiscreteScheduler - scheduler_classes = [FlowMatchEulerDiscreteScheduler] scheduler_kwargs = {} transformer_kwargs = { @@ -124,30 +123,29 @@ class CogView4LoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): """ Tests a simple usecase where users could use saving utilities for LoRA through save_pretrained """ - for scheduler_cls in self.scheduler_classes: - components, _, _ = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + components, _, _ = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue(output_no_lora.shape == self.output_shape) - images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - with tempfile.TemporaryDirectory() as tmpdirname: - pipe.save_pretrained(tmpdirname) + with tempfile.TemporaryDirectory() as tmpdirname: + pipe.save_pretrained(tmpdirname) - pipe_from_pretrained = self.pipeline_class.from_pretrained(tmpdirname) - pipe_from_pretrained.to(torch_device) + pipe_from_pretrained = self.pipeline_class.from_pretrained(tmpdirname) + pipe_from_pretrained.to(torch_device) - images_lora_save_pretrained = pipe_from_pretrained(**inputs, generator=torch.manual_seed(0))[0] + images_lora_save_pretrained = pipe_from_pretrained(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - np.allclose(images_lora, images_lora_save_pretrained, atol=1e-3, rtol=1e-3), - "Loading from saved checkpoints should give same results.", - ) + self.assertTrue( + np.allclose(images_lora, images_lora_save_pretrained, atol=1e-3, rtol=1e-3), + "Loading from saved checkpoints should give same results.", + ) @parameterized.expand([("block_level", True), ("leaf_level", False)]) @require_torch_accelerator diff --git a/tests/lora/test_lora_layers_flux.py b/tests/lora/test_lora_layers_flux.py index e6048f509f..6c22a34889 100644 --- a/tests/lora/test_lora_layers_flux.py +++ b/tests/lora/test_lora_layers_flux.py @@ -55,9 +55,8 @@ from .utils import PeftLoraLoaderMixinTests, check_if_lora_correctly_set # noqa @require_peft_backend class FluxLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): pipeline_class = FluxPipeline - scheduler_cls = FlowMatchEulerDiscreteScheduler() + scheduler_cls = FlowMatchEulerDiscreteScheduler scheduler_kwargs = {} - scheduler_classes = [FlowMatchEulerDiscreteScheduler] transformer_kwargs = { "patch_size": 1, "in_channels": 4, @@ -282,9 +281,8 @@ class FluxLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): class FluxControlLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): pipeline_class = FluxControlPipeline - scheduler_cls = FlowMatchEulerDiscreteScheduler() + scheduler_cls = FlowMatchEulerDiscreteScheduler scheduler_kwargs = {} - scheduler_classes = [FlowMatchEulerDiscreteScheduler] transformer_kwargs = { "patch_size": 1, "in_channels": 8, diff --git a/tests/lora/test_lora_layers_hunyuanvideo.py b/tests/lora/test_lora_layers_hunyuanvideo.py index 62d045f836..7ea0f1fcc9 100644 --- a/tests/lora/test_lora_layers_hunyuanvideo.py +++ b/tests/lora/test_lora_layers_hunyuanvideo.py @@ -51,7 +51,6 @@ from .utils import PeftLoraLoaderMixinTests # noqa: E402 class HunyuanVideoLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): pipeline_class = HunyuanVideoPipeline scheduler_cls = FlowMatchEulerDiscreteScheduler - scheduler_classes = [FlowMatchEulerDiscreteScheduler] scheduler_kwargs = {} transformer_kwargs = { diff --git a/tests/lora/test_lora_layers_ltx_video.py b/tests/lora/test_lora_layers_ltx_video.py index a8ad30e448..6ab51a5e51 100644 --- a/tests/lora/test_lora_layers_ltx_video.py +++ b/tests/lora/test_lora_layers_ltx_video.py @@ -37,7 +37,6 @@ from .utils import PeftLoraLoaderMixinTests # noqa: E402 class LTXVideoLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): pipeline_class = LTXPipeline scheduler_cls = FlowMatchEulerDiscreteScheduler - scheduler_classes = [FlowMatchEulerDiscreteScheduler] scheduler_kwargs = {} transformer_kwargs = { diff --git a/tests/lora/test_lora_layers_lumina2.py b/tests/lora/test_lora_layers_lumina2.py index 0ebc831b11..0417b05b33 100644 --- a/tests/lora/test_lora_layers_lumina2.py +++ b/tests/lora/test_lora_layers_lumina2.py @@ -39,7 +39,6 @@ from .utils import PeftLoraLoaderMixinTests, check_if_lora_correctly_set # noqa class Lumina2LoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): pipeline_class = Lumina2Pipeline scheduler_cls = FlowMatchEulerDiscreteScheduler - scheduler_classes = [FlowMatchEulerDiscreteScheduler] scheduler_kwargs = {} transformer_kwargs = { @@ -141,33 +140,30 @@ class Lumina2LoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): strict=False, ) def test_lora_fuse_nan(self): - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder" - ) + if "text_encoder" in self.pipeline_class._lora_loadable_modules: + pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") + self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") + denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet + denoiser.add_adapter(denoiser_lora_config, "adapter-1") + self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - # corrupt one LoRA weight with `inf` values - with torch.no_grad(): - pipe.transformer.layers[0].attn.to_q.lora_A["adapter-1"].weight += float("inf") + # corrupt one LoRA weight with `inf` values + with torch.no_grad(): + pipe.transformer.layers[0].attn.to_q.lora_A["adapter-1"].weight += float("inf") - # with `safe_fusing=True` we should see an Error - with self.assertRaises(ValueError): - pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules, safe_fusing=True) + # with `safe_fusing=True` we should see an Error + with self.assertRaises(ValueError): + pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules, safe_fusing=True) - # without we should not see an error, but every image will be black - pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules, safe_fusing=False) - out = pipe(**inputs)[0] + # without we should not see an error, but every image will be black + pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules, safe_fusing=False) + out = pipe(**inputs)[0] - self.assertTrue(np.isnan(out).all()) + self.assertTrue(np.isnan(out).all()) diff --git a/tests/lora/test_lora_layers_mochi.py b/tests/lora/test_lora_layers_mochi.py index 21cc5f11a3..7be81273db 100644 --- a/tests/lora/test_lora_layers_mochi.py +++ b/tests/lora/test_lora_layers_mochi.py @@ -37,7 +37,6 @@ from .utils import PeftLoraLoaderMixinTests # noqa: E402 class MochiLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): pipeline_class = MochiPipeline scheduler_cls = FlowMatchEulerDiscreteScheduler - scheduler_classes = [FlowMatchEulerDiscreteScheduler] scheduler_kwargs = {} transformer_kwargs = { diff --git a/tests/lora/test_lora_layers_qwenimage.py b/tests/lora/test_lora_layers_qwenimage.py index 44ef9b0a37..51de2f8e20 100644 --- a/tests/lora/test_lora_layers_qwenimage.py +++ b/tests/lora/test_lora_layers_qwenimage.py @@ -37,7 +37,6 @@ from .utils import PeftLoraLoaderMixinTests # noqa: E402 class QwenImageLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): pipeline_class = QwenImagePipeline scheduler_cls = FlowMatchEulerDiscreteScheduler - scheduler_classes = [FlowMatchEulerDiscreteScheduler] scheduler_kwargs = {} transformer_kwargs = { diff --git a/tests/lora/test_lora_layers_sana.py b/tests/lora/test_lora_layers_sana.py index a08908c610..3cdb28de75 100644 --- a/tests/lora/test_lora_layers_sana.py +++ b/tests/lora/test_lora_layers_sana.py @@ -31,9 +31,8 @@ from .utils import PeftLoraLoaderMixinTests # noqa: E402 @require_peft_backend class SanaLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): pipeline_class = SanaPipeline - scheduler_cls = FlowMatchEulerDiscreteScheduler(shift=7.0) - scheduler_kwargs = {} - scheduler_classes = [FlowMatchEulerDiscreteScheduler] + scheduler_cls = FlowMatchEulerDiscreteScheduler + scheduler_kwargs = {"shift": 7.0} transformer_kwargs = { "patch_size": 1, "in_channels": 4, diff --git a/tests/lora/test_lora_layers_sd3.py b/tests/lora/test_lora_layers_sd3.py index 95f6f325e4..228460eaad 100644 --- a/tests/lora/test_lora_layers_sd3.py +++ b/tests/lora/test_lora_layers_sd3.py @@ -55,7 +55,6 @@ class SD3LoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): pipeline_class = StableDiffusion3Pipeline scheduler_cls = FlowMatchEulerDiscreteScheduler scheduler_kwargs = {} - scheduler_classes = [FlowMatchEulerDiscreteScheduler] transformer_kwargs = { "sample_size": 32, "patch_size": 1, diff --git a/tests/lora/test_lora_layers_wan.py b/tests/lora/test_lora_layers_wan.py index 0ba80d2be1..5734509b41 100644 --- a/tests/lora/test_lora_layers_wan.py +++ b/tests/lora/test_lora_layers_wan.py @@ -42,7 +42,6 @@ from .utils import PeftLoraLoaderMixinTests # noqa: E402 class WanLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): pipeline_class = WanPipeline scheduler_cls = FlowMatchEulerDiscreteScheduler - scheduler_classes = [FlowMatchEulerDiscreteScheduler] scheduler_kwargs = {} transformer_kwargs = { diff --git a/tests/lora/test_lora_layers_wanvace.py b/tests/lora/test_lora_layers_wanvace.py index d8dde32dd8..c3244e150e 100644 --- a/tests/lora/test_lora_layers_wanvace.py +++ b/tests/lora/test_lora_layers_wanvace.py @@ -50,7 +50,6 @@ from .utils import PeftLoraLoaderMixinTests # noqa: E402 class WanVACELoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): pipeline_class = WanVACEPipeline scheduler_cls = FlowMatchEulerDiscreteScheduler - scheduler_classes = [FlowMatchEulerDiscreteScheduler] scheduler_kwargs = {} transformer_kwargs = { @@ -165,9 +164,8 @@ class WanVACELoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): @require_peft_version_greater("0.13.2") def test_lora_exclude_modules_wanvace(self): - scheduler_cls = self.scheduler_classes[0] exclude_module_name = "vace_blocks.0.proj_out" - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() pipe = self.pipeline_class(**components).to(torch_device) _, _, inputs = self.get_dummy_inputs(with_generator=False) diff --git a/tests/lora/utils.py b/tests/lora/utils.py index 72c1dddaa2..ecaa553ce4 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -26,8 +26,6 @@ from parameterized import parameterized from diffusers import ( AutoencoderKL, - DDIMScheduler, - LCMScheduler, UNet2DConditionModel, ) from diffusers.utils import logging @@ -109,7 +107,6 @@ class PeftLoraLoaderMixinTests: scheduler_cls = None scheduler_kwargs = None - scheduler_classes = [DDIMScheduler, LCMScheduler] has_two_text_encoders = False has_three_text_encoders = False @@ -129,13 +126,13 @@ class PeftLoraLoaderMixinTests: text_encoder_target_modules = ["q_proj", "k_proj", "v_proj", "out_proj"] denoiser_target_modules = ["to_q", "to_k", "to_v", "to_out.0"] - def get_dummy_components(self, scheduler_cls=None, use_dora=False, lora_alpha=None): + def get_dummy_components(self, use_dora=False, lora_alpha=None): if self.unet_kwargs and self.transformer_kwargs: raise ValueError("Both `unet_kwargs` and `transformer_kwargs` cannot be specified.") if self.has_two_text_encoders and self.has_three_text_encoders: raise ValueError("Both `has_two_text_encoders` and `has_three_text_encoders` cannot be True.") - scheduler_cls = self.scheduler_cls if scheduler_cls is None else scheduler_cls + scheduler_cls = self.scheduler_cls rank = 4 lora_alpha = rank if lora_alpha is None else lora_alpha @@ -319,152 +316,143 @@ class PeftLoraLoaderMixinTests: """ Tests a simple inference and makes sure it works as expected """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, _ = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) + components, text_lora_config, _ = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs() - output_no_lora = pipe(**inputs)[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + _, _, inputs = self.get_dummy_inputs() + output_no_lora = pipe(**inputs)[0] + self.assertTrue(output_no_lora.shape == self.output_shape) def test_simple_inference_with_text_lora(self): """ Tests a simple inference with lora attached on the text encoder and makes sure it works as expected """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, _ = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + components, text_lora_config, _ = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue(output_no_lora.shape == self.output_shape) - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) + pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) - output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - not np.allclose(output_lora, output_no_lora, atol=1e-3, rtol=1e-3), "Lora should change the output" - ) + output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue( + not np.allclose(output_lora, output_no_lora, atol=1e-3, rtol=1e-3), "Lora should change the output" + ) @require_peft_version_greater("0.13.1") def test_low_cpu_mem_usage_with_injection(self): """Tests if we can inject LoRA state dict with low_cpu_mem_usage.""" - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - inject_adapter_in_model(text_lora_config, pipe.text_encoder, low_cpu_mem_usage=True) + if "text_encoder" in self.pipeline_class._lora_loadable_modules: + inject_adapter_in_model(text_lora_config, pipe.text_encoder, low_cpu_mem_usage=True) + self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder.") + self.assertTrue( + "meta" in {p.device.type for p in pipe.text_encoder.parameters()}, + "The LoRA params should be on 'meta' device.", + ) + + te_state_dict = initialize_dummy_state_dict(get_peft_model_state_dict(pipe.text_encoder)) + set_peft_model_state_dict(pipe.text_encoder, te_state_dict, low_cpu_mem_usage=True) + self.assertTrue( + "meta" not in {p.device.type for p in pipe.text_encoder.parameters()}, + "No param should be on 'meta' device.", + ) + + denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet + inject_adapter_in_model(denoiser_lora_config, denoiser, low_cpu_mem_usage=True) + self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") + self.assertTrue( + "meta" in {p.device.type for p in denoiser.parameters()}, "The LoRA params should be on 'meta' device." + ) + + denoiser_state_dict = initialize_dummy_state_dict(get_peft_model_state_dict(denoiser)) + set_peft_model_state_dict(denoiser, denoiser_state_dict, low_cpu_mem_usage=True) + self.assertTrue( + "meta" not in {p.device.type for p in denoiser.parameters()}, "No param should be on 'meta' device." + ) + + if self.has_two_text_encoders or self.has_three_text_encoders: + if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: + inject_adapter_in_model(text_lora_config, pipe.text_encoder_2, low_cpu_mem_usage=True) self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder." + check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" ) self.assertTrue( - "meta" in {p.device.type for p in pipe.text_encoder.parameters()}, + "meta" in {p.device.type for p in pipe.text_encoder_2.parameters()}, "The LoRA params should be on 'meta' device.", ) - te_state_dict = initialize_dummy_state_dict(get_peft_model_state_dict(pipe.text_encoder)) - set_peft_model_state_dict(pipe.text_encoder, te_state_dict, low_cpu_mem_usage=True) + te2_state_dict = initialize_dummy_state_dict(get_peft_model_state_dict(pipe.text_encoder_2)) + set_peft_model_state_dict(pipe.text_encoder_2, te2_state_dict, low_cpu_mem_usage=True) self.assertTrue( - "meta" not in {p.device.type for p in pipe.text_encoder.parameters()}, + "meta" not in {p.device.type for p in pipe.text_encoder_2.parameters()}, "No param should be on 'meta' device.", ) - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - inject_adapter_in_model(denoiser_lora_config, denoiser, low_cpu_mem_usage=True) - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - self.assertTrue( - "meta" in {p.device.type for p in denoiser.parameters()}, "The LoRA params should be on 'meta' device." - ) - - denoiser_state_dict = initialize_dummy_state_dict(get_peft_model_state_dict(denoiser)) - set_peft_model_state_dict(denoiser, denoiser_state_dict, low_cpu_mem_usage=True) - self.assertTrue( - "meta" not in {p.device.type for p in denoiser.parameters()}, "No param should be on 'meta' device." - ) - - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - inject_adapter_in_model(text_lora_config, pipe.text_encoder_2, low_cpu_mem_usage=True) - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) - self.assertTrue( - "meta" in {p.device.type for p in pipe.text_encoder_2.parameters()}, - "The LoRA params should be on 'meta' device.", - ) - - te2_state_dict = initialize_dummy_state_dict(get_peft_model_state_dict(pipe.text_encoder_2)) - set_peft_model_state_dict(pipe.text_encoder_2, te2_state_dict, low_cpu_mem_usage=True) - self.assertTrue( - "meta" not in {p.device.type for p in pipe.text_encoder_2.parameters()}, - "No param should be on 'meta' device.", - ) - - _, _, inputs = self.get_dummy_inputs() - output_lora = pipe(**inputs)[0] - self.assertTrue(output_lora.shape == self.output_shape) + _, _, inputs = self.get_dummy_inputs() + output_lora = pipe(**inputs)[0] + self.assertTrue(output_lora.shape == self.output_shape) @require_peft_version_greater("0.13.1") @require_transformers_version_greater("4.45.2") def test_low_cpu_mem_usage_with_loading(self): """Tests if we can load LoRA state dict with low_cpu_mem_usage.""" + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue(output_no_lora.shape == self.output_shape) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) + images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + with tempfile.TemporaryDirectory() as tmpdirname: + modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True) + lora_state_dicts = self._get_lora_state_dicts(modules_to_save) + self.pipeline_class.save_lora_weights( + save_directory=tmpdirname, safe_serialization=False, **lora_state_dicts + ) - with tempfile.TemporaryDirectory() as tmpdirname: - modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True) - lora_state_dicts = self._get_lora_state_dicts(modules_to_save) - self.pipeline_class.save_lora_weights( - save_directory=tmpdirname, safe_serialization=False, **lora_state_dicts - ) + self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.bin"))) + pipe.unload_lora_weights() + pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin"), low_cpu_mem_usage=False) - self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.bin"))) - pipe.unload_lora_weights() - pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin"), low_cpu_mem_usage=False) + for module_name, module in modules_to_save.items(): + self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") - for module_name, module in modules_to_save.items(): - self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") + images_lora_from_pretrained = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue( + np.allclose(images_lora, images_lora_from_pretrained, atol=1e-3, rtol=1e-3), + "Loading from saved checkpoints should give same results.", + ) - images_lora_from_pretrained = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - np.allclose(images_lora, images_lora_from_pretrained, atol=1e-3, rtol=1e-3), - "Loading from saved checkpoints should give same results.", - ) + # Now, check for `low_cpu_mem_usage.` + pipe.unload_lora_weights() + pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin"), low_cpu_mem_usage=True) - # Now, check for `low_cpu_mem_usage.` - pipe.unload_lora_weights() - pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin"), low_cpu_mem_usage=True) + for module_name, module in modules_to_save.items(): + self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") - for module_name, module in modules_to_save.items(): - self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") - - images_lora_from_pretrained_low_cpu = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - np.allclose( - images_lora_from_pretrained_low_cpu, images_lora_from_pretrained, atol=1e-3, rtol=1e-3 - ), - "Loading from saved checkpoints with `low_cpu_mem_usage` should give same results.", - ) + images_lora_from_pretrained_low_cpu = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue( + np.allclose(images_lora_from_pretrained_low_cpu, images_lora_from_pretrained, atol=1e-3, rtol=1e-3), + "Loading from saved checkpoints with `low_cpu_mem_usage` should give same results.", + ) def test_simple_inference_with_text_lora_and_scale(self): """ @@ -472,411 +460,393 @@ class PeftLoraLoaderMixinTests: and makes sure it works as expected """ attention_kwargs_name = determine_attention_kwargs_name(self.pipeline_class) + components, text_lora_config, _ = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, _ = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue(output_no_lora.shape == self.output_shape) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) + output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue( + not np.allclose(output_lora, output_no_lora, atol=1e-3, rtol=1e-3), "Lora should change the output" + ) - output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - not np.allclose(output_lora, output_no_lora, atol=1e-3, rtol=1e-3), "Lora should change the output" - ) + attention_kwargs = {attention_kwargs_name: {"scale": 0.5}} + output_lora_scale = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] - attention_kwargs = {attention_kwargs_name: {"scale": 0.5}} - output_lora_scale = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] + self.assertTrue( + not np.allclose(output_lora, output_lora_scale, atol=1e-3, rtol=1e-3), + "Lora + scale should change the output", + ) - self.assertTrue( - not np.allclose(output_lora, output_lora_scale, atol=1e-3, rtol=1e-3), - "Lora + scale should change the output", - ) + attention_kwargs = {attention_kwargs_name: {"scale": 0.0}} + output_lora_0_scale = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] - attention_kwargs = {attention_kwargs_name: {"scale": 0.0}} - output_lora_0_scale = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] - - self.assertTrue( - np.allclose(output_no_lora, output_lora_0_scale, atol=1e-3, rtol=1e-3), - "Lora + 0 scale should lead to same result as no LoRA", - ) + self.assertTrue( + np.allclose(output_no_lora, output_lora_0_scale, atol=1e-3, rtol=1e-3), + "Lora + 0 scale should lead to same result as no LoRA", + ) def test_simple_inference_with_text_lora_fused(self): """ Tests a simple inference with lora attached into text encoder + fuses the lora weights into base model and makes sure it works as expected """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, _ = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + components, text_lora_config, _ = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue(output_no_lora.shape == self.output_shape) - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) + pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) - pipe.fuse_lora() - # Fusing should still keep the LoRA layers - self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") + pipe.fuse_lora() + # Fusing should still keep the LoRA layers + self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) + if self.has_two_text_encoders or self.has_three_text_encoders: + if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: + self.assertTrue( + check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" + ) - ouput_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse( - np.allclose(ouput_fused, output_no_lora, atol=1e-3, rtol=1e-3), "Fused lora should change the output" - ) + ouput_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertFalse( + np.allclose(ouput_fused, output_no_lora, atol=1e-3, rtol=1e-3), "Fused lora should change the output" + ) def test_simple_inference_with_text_lora_unloaded(self): """ Tests a simple inference with lora attached to text encoder, then unloads the lora weights and makes sure it works as expected """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, _ = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + components, text_lora_config, _ = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue(output_no_lora.shape == self.output_shape) - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) + pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) - pipe.unload_lora_weights() - # unloading should remove the LoRA layers - self.assertFalse( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly unloaded in text encoder" - ) + pipe.unload_lora_weights() + # unloading should remove the LoRA layers + self.assertFalse(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly unloaded in text encoder") - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - self.assertFalse( - check_if_lora_correctly_set(pipe.text_encoder_2), - "Lora not correctly unloaded in text encoder 2", - ) + if self.has_two_text_encoders or self.has_three_text_encoders: + if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: + self.assertFalse( + check_if_lora_correctly_set(pipe.text_encoder_2), + "Lora not correctly unloaded in text encoder 2", + ) - ouput_unloaded = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - np.allclose(ouput_unloaded, output_no_lora, atol=1e-3, rtol=1e-3), - "Fused lora should change the output", - ) + ouput_unloaded = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue( + np.allclose(ouput_unloaded, output_no_lora, atol=1e-3, rtol=1e-3), + "Fused lora should change the output", + ) def test_simple_inference_with_text_lora_save_load(self): """ Tests a simple usecase where users could use saving utilities for LoRA. """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, _ = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + components, text_lora_config, _ = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue(output_no_lora.shape == self.output_shape) - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) + pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) - images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - with tempfile.TemporaryDirectory() as tmpdirname: - modules_to_save = self._get_modules_to_save(pipe) - lora_state_dicts = self._get_lora_state_dicts(modules_to_save) + with tempfile.TemporaryDirectory() as tmpdirname: + modules_to_save = self._get_modules_to_save(pipe) + lora_state_dicts = self._get_lora_state_dicts(modules_to_save) - self.pipeline_class.save_lora_weights( - save_directory=tmpdirname, safe_serialization=False, **lora_state_dicts - ) - - self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.bin"))) - pipe.unload_lora_weights() - pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin")) - - for module_name, module in modules_to_save.items(): - self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") - - images_lora_from_pretrained = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(images_lora, images_lora_from_pretrained, atol=1e-3, rtol=1e-3), - "Loading from saved checkpoints should give same results.", + self.pipeline_class.save_lora_weights( + save_directory=tmpdirname, safe_serialization=False, **lora_state_dicts ) + self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.bin"))) + pipe.unload_lora_weights() + pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin")) + + for module_name, module in modules_to_save.items(): + self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") + + images_lora_from_pretrained = pipe(**inputs, generator=torch.manual_seed(0))[0] + + self.assertTrue( + np.allclose(images_lora, images_lora_from_pretrained, atol=1e-3, rtol=1e-3), + "Loading from saved checkpoints should give same results.", + ) + def test_simple_inference_with_partial_text_lora(self): """ Tests a simple inference with lora attached on the text encoder with different ranks and some adapters removed and makes sure it works as expected """ - for scheduler_cls in self.scheduler_classes: - components, _, _ = self.get_dummy_components(scheduler_cls) - # Verify `StableDiffusionLoraLoaderMixin.load_lora_into_text_encoder` handles different ranks per module (PR#8324). - text_lora_config = LoraConfig( - r=4, - rank_pattern={self.text_encoder_target_modules[i]: i + 1 for i in range(3)}, - lora_alpha=4, - target_modules=self.text_encoder_target_modules, - init_lora_weights=False, - use_dora=False, - ) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + components, _, _ = self.get_dummy_components() + # Verify `StableDiffusionLoraLoaderMixin.load_lora_into_text_encoder` handles different ranks per module (PR#8324). + text_lora_config = LoraConfig( + r=4, + rank_pattern={self.text_encoder_target_modules[i]: i + 1 for i in range(3)}, + lora_alpha=4, + target_modules=self.text_encoder_target_modules, + init_lora_weights=False, + use_dora=False, + ) + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue(output_no_lora.shape == self.output_shape) - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) + pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) - state_dict = {} - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - # Gather the state dict for the PEFT model, excluding `layers.4`, to ensure `load_lora_into_text_encoder` - # supports missing layers (PR#8324). - state_dict = { - f"text_encoder.{module_name}": param - for module_name, param in get_peft_model_state_dict(pipe.text_encoder).items() - if "text_model.encoder.layers.4" not in module_name - } + state_dict = {} + if "text_encoder" in self.pipeline_class._lora_loadable_modules: + # Gather the state dict for the PEFT model, excluding `layers.4`, to ensure `load_lora_into_text_encoder` + # supports missing layers (PR#8324). + state_dict = { + f"text_encoder.{module_name}": param + for module_name, param in get_peft_model_state_dict(pipe.text_encoder).items() + if "text_model.encoder.layers.4" not in module_name + } - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - state_dict.update( - { - f"text_encoder_2.{module_name}": param - for module_name, param in get_peft_model_state_dict(pipe.text_encoder_2).items() - if "text_model.encoder.layers.4" not in module_name - } - ) + if self.has_two_text_encoders or self.has_three_text_encoders: + if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: + state_dict.update( + { + f"text_encoder_2.{module_name}": param + for module_name, param in get_peft_model_state_dict(pipe.text_encoder_2).items() + if "text_model.encoder.layers.4" not in module_name + } + ) - output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - not np.allclose(output_lora, output_no_lora, atol=1e-3, rtol=1e-3), "Lora should change the output" - ) + output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue( + not np.allclose(output_lora, output_no_lora, atol=1e-3, rtol=1e-3), "Lora should change the output" + ) - # Unload lora and load it back using the pipe.load_lora_weights machinery - pipe.unload_lora_weights() - pipe.load_lora_weights(state_dict) + # Unload lora and load it back using the pipe.load_lora_weights machinery + pipe.unload_lora_weights() + pipe.load_lora_weights(state_dict) - output_partial_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - not np.allclose(output_partial_lora, output_lora, atol=1e-3, rtol=1e-3), - "Removing adapters should change the output", - ) + output_partial_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue( + not np.allclose(output_partial_lora, output_lora, atol=1e-3, rtol=1e-3), + "Removing adapters should change the output", + ) def test_simple_inference_save_pretrained_with_text_lora(self): """ Tests a simple usecase where users could use saving utilities for LoRA through save_pretrained """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, _ = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + components, text_lora_config, _ = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue(output_no_lora.shape == self.output_shape) - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) - images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) + images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - with tempfile.TemporaryDirectory() as tmpdirname: - pipe.save_pretrained(tmpdirname) + with tempfile.TemporaryDirectory() as tmpdirname: + pipe.save_pretrained(tmpdirname) - pipe_from_pretrained = self.pipeline_class.from_pretrained(tmpdirname) - pipe_from_pretrained.to(torch_device) + pipe_from_pretrained = self.pipeline_class.from_pretrained(tmpdirname) + pipe_from_pretrained.to(torch_device) - if "text_encoder" in self.pipeline_class._lora_loadable_modules: + if "text_encoder" in self.pipeline_class._lora_loadable_modules: + self.assertTrue( + check_if_lora_correctly_set(pipe_from_pretrained.text_encoder), + "Lora not correctly set in text encoder", + ) + + if self.has_two_text_encoders or self.has_three_text_encoders: + if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: self.assertTrue( - check_if_lora_correctly_set(pipe_from_pretrained.text_encoder), - "Lora not correctly set in text encoder", + check_if_lora_correctly_set(pipe_from_pretrained.text_encoder_2), + "Lora not correctly set in text encoder 2", ) - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - self.assertTrue( - check_if_lora_correctly_set(pipe_from_pretrained.text_encoder_2), - "Lora not correctly set in text encoder 2", - ) + images_lora_save_pretrained = pipe_from_pretrained(**inputs, generator=torch.manual_seed(0))[0] - images_lora_save_pretrained = pipe_from_pretrained(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(images_lora, images_lora_save_pretrained, atol=1e-3, rtol=1e-3), - "Loading from saved checkpoints should give same results.", - ) + self.assertTrue( + np.allclose(images_lora, images_lora_save_pretrained, atol=1e-3, rtol=1e-3), + "Loading from saved checkpoints should give same results.", + ) def test_simple_inference_with_text_denoiser_lora_save_load(self): """ Tests a simple usecase where users could use saving utilities for LoRA for Unet + text encoder """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue(output_no_lora.shape == self.output_shape) - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) + pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - with tempfile.TemporaryDirectory() as tmpdirname: - modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True) - lora_state_dicts = self._get_lora_state_dicts(modules_to_save) - self.pipeline_class.save_lora_weights( - save_directory=tmpdirname, safe_serialization=False, **lora_state_dicts - ) - - self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.bin"))) - pipe.unload_lora_weights() - pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin")) - - for module_name, module in modules_to_save.items(): - self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") - - images_lora_from_pretrained = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - np.allclose(images_lora, images_lora_from_pretrained, atol=1e-3, rtol=1e-3), - "Loading from saved checkpoints should give same results.", + with tempfile.TemporaryDirectory() as tmpdirname: + modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True) + lora_state_dicts = self._get_lora_state_dicts(modules_to_save) + self.pipeline_class.save_lora_weights( + save_directory=tmpdirname, safe_serialization=False, **lora_state_dicts ) + self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.bin"))) + pipe.unload_lora_weights() + pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin")) + + for module_name, module in modules_to_save.items(): + self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") + + images_lora_from_pretrained = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue( + np.allclose(images_lora, images_lora_from_pretrained, atol=1e-3, rtol=1e-3), + "Loading from saved checkpoints should give same results.", + ) + def test_simple_inference_with_text_denoiser_lora_and_scale(self): """ Tests a simple inference with lora attached on the text encoder + Unet + scale argument and makes sure it works as expected """ attention_kwargs_name = determine_attention_kwargs_name(self.pipeline_class) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue(output_no_lora.shape == self.output_shape) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) + output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue( + not np.allclose(output_lora, output_no_lora, atol=1e-3, rtol=1e-3), "Lora should change the output" + ) - output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + attention_kwargs = {attention_kwargs_name: {"scale": 0.5}} + output_lora_scale = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] + + self.assertTrue( + not np.allclose(output_lora, output_lora_scale, atol=1e-3, rtol=1e-3), + "Lora + scale should change the output", + ) + + attention_kwargs = {attention_kwargs_name: {"scale": 0.0}} + output_lora_0_scale = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] + + self.assertTrue( + np.allclose(output_no_lora, output_lora_0_scale, atol=1e-3, rtol=1e-3), + "Lora + 0 scale should lead to same result as no LoRA", + ) + + if "text_encoder" in self.pipeline_class._lora_loadable_modules: self.assertTrue( - not np.allclose(output_lora, output_no_lora, atol=1e-3, rtol=1e-3), "Lora should change the output" + pipe.text_encoder.text_model.encoder.layers[0].self_attn.q_proj.scaling["default"] == 1.0, + "The scaling parameter has not been correctly restored!", ) - attention_kwargs = {attention_kwargs_name: {"scale": 0.5}} - output_lora_scale = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] - - self.assertTrue( - not np.allclose(output_lora, output_lora_scale, atol=1e-3, rtol=1e-3), - "Lora + scale should change the output", - ) - - attention_kwargs = {attention_kwargs_name: {"scale": 0.0}} - output_lora_0_scale = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] - - self.assertTrue( - np.allclose(output_no_lora, output_lora_0_scale, atol=1e-3, rtol=1e-3), - "Lora + 0 scale should lead to same result as no LoRA", - ) - - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - self.assertTrue( - pipe.text_encoder.text_model.encoder.layers[0].self_attn.q_proj.scaling["default"] == 1.0, - "The scaling parameter has not been correctly restored!", - ) - def test_simple_inference_with_text_lora_denoiser_fused(self): """ Tests a simple inference with lora attached into text encoder + fuses the lora weights into base model and makes sure it works as expected - with unet """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue(output_no_lora.shape == self.output_shape) - pipe, denoiser = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) + pipe, denoiser = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules) + pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules) - # Fusing should still keep the LoRA layers - if "text_encoder" in self.pipeline_class._lora_loadable_modules: + # Fusing should still keep the LoRA layers + if "text_encoder" in self.pipeline_class._lora_loadable_modules: + self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") + + self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser") + + if self.has_two_text_encoders or self.has_three_text_encoders: + if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder" + check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" ) - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser") - - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) - - output_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse( - np.allclose(output_fused, output_no_lora, atol=1e-3, rtol=1e-3), "Fused lora should change the output" - ) + output_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertFalse( + np.allclose(output_fused, output_no_lora, atol=1e-3, rtol=1e-3), "Fused lora should change the output" + ) def test_simple_inference_with_text_denoiser_lora_unloaded(self): """ Tests a simple inference with lora attached to text encoder and unet, then unloads the lora weights and makes sure it works as expected """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue(output_no_lora.shape == self.output_shape) - pipe, denoiser = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) + pipe, denoiser = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - pipe.unload_lora_weights() - # unloading should remove the LoRA layers - self.assertFalse( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly unloaded in text encoder" - ) - self.assertFalse(check_if_lora_correctly_set(denoiser), "Lora not correctly unloaded in denoiser") + pipe.unload_lora_weights() + # unloading should remove the LoRA layers + self.assertFalse(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly unloaded in text encoder") + self.assertFalse(check_if_lora_correctly_set(denoiser), "Lora not correctly unloaded in denoiser") - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - self.assertFalse( - check_if_lora_correctly_set(pipe.text_encoder_2), - "Lora not correctly unloaded in text encoder 2", - ) + if self.has_two_text_encoders or self.has_three_text_encoders: + if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: + self.assertFalse( + check_if_lora_correctly_set(pipe.text_encoder_2), + "Lora not correctly unloaded in text encoder 2", + ) - output_unloaded = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - np.allclose(output_unloaded, output_no_lora, atol=1e-3, rtol=1e-3), - "Fused lora should change the output", - ) + output_unloaded = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue( + np.allclose(output_unloaded, output_no_lora, atol=1e-3, rtol=1e-3), + "Fused lora should change the output", + ) def test_simple_inference_with_text_denoiser_lora_unfused( self, expected_atol: float = 1e-3, expected_rtol: float = 1e-3 @@ -885,125 +855,120 @@ class PeftLoraLoaderMixinTests: Tests a simple inference with lora attached to text encoder and unet, then unloads the lora weights and makes sure it works as expected """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) - pipe, denoiser = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) + pipe, denoiser = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules) - self.assertTrue(pipe.num_fused_loras == 1, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") - output_fused_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules) + self.assertTrue(pipe.num_fused_loras == 1, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") + output_fused_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - pipe.unfuse_lora(components=self.pipeline_class._lora_loadable_modules) - self.assertTrue(pipe.num_fused_loras == 0, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") - output_unfused_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + pipe.unfuse_lora(components=self.pipeline_class._lora_loadable_modules) + self.assertTrue(pipe.num_fused_loras == 0, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") + output_unfused_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - # unloading should remove the LoRA layers - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Unfuse should still keep LoRA layers") + # unloading should remove the LoRA layers + if "text_encoder" in self.pipeline_class._lora_loadable_modules: + self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Unfuse should still keep LoRA layers") - self.assertTrue(check_if_lora_correctly_set(denoiser), "Unfuse should still keep LoRA layers") + self.assertTrue(check_if_lora_correctly_set(denoiser), "Unfuse should still keep LoRA layers") - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Unfuse should still keep LoRA layers" - ) + if self.has_two_text_encoders or self.has_three_text_encoders: + if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: + self.assertTrue( + check_if_lora_correctly_set(pipe.text_encoder_2), "Unfuse should still keep LoRA layers" + ) - # Fuse and unfuse should lead to the same results - self.assertTrue( - np.allclose(output_fused_lora, output_unfused_lora, atol=expected_atol, rtol=expected_rtol), - "Fused lora should not change the output", - ) + # Fuse and unfuse should lead to the same results + self.assertTrue( + np.allclose(output_fused_lora, output_unfused_lora, atol=expected_atol, rtol=expected_rtol), + "Fused lora should not change the output", + ) def test_simple_inference_with_text_denoiser_multi_adapter(self): """ Tests a simple inference with lora attached to text encoder and unet, attaches multiple adapters and set them """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") + if "text_encoder" in self.pipeline_class._lora_loadable_modules: + pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") + pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") + self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") + + denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet + denoiser.add_adapter(denoiser_lora_config, "adapter-1") + denoiser.add_adapter(denoiser_lora_config, "adapter-2") + self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") + + if self.has_two_text_encoders or self.has_three_text_encoders: + if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: + pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") + pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-2") self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder" + check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" ) - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") - denoiser.add_adapter(denoiser_lora_config, "adapter-2") - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") + pipe.set_adapters("adapter-1") + output_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertFalse( + np.allclose(output_no_lora, output_adapter_1, atol=1e-3, rtol=1e-3), + "Adapter outputs should be different.", + ) - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-2") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) + pipe.set_adapters("adapter-2") + output_adapter_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertFalse( + np.allclose(output_no_lora, output_adapter_2, atol=1e-3, rtol=1e-3), + "Adapter outputs should be different.", + ) - pipe.set_adapters("adapter-1") - output_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse( - np.allclose(output_no_lora, output_adapter_1, atol=1e-3, rtol=1e-3), - "Adapter outputs should be different.", - ) + pipe.set_adapters(["adapter-1", "adapter-2"]) + output_adapter_mixed = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertFalse( + np.allclose(output_no_lora, output_adapter_mixed, atol=1e-3, rtol=1e-3), + "Adapter outputs should be different.", + ) - pipe.set_adapters("adapter-2") - output_adapter_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse( - np.allclose(output_no_lora, output_adapter_2, atol=1e-3, rtol=1e-3), - "Adapter outputs should be different.", - ) + # Fuse and unfuse should lead to the same results + self.assertFalse( + np.allclose(output_adapter_1, output_adapter_2, atol=1e-3, rtol=1e-3), + "Adapter 1 and 2 should give different results", + ) - pipe.set_adapters(["adapter-1", "adapter-2"]) - output_adapter_mixed = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse( - np.allclose(output_no_lora, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter outputs should be different.", - ) + self.assertFalse( + np.allclose(output_adapter_1, output_adapter_mixed, atol=1e-3, rtol=1e-3), + "Adapter 1 and mixed adapters should give different results", + ) - # Fuse and unfuse should lead to the same results - self.assertFalse( - np.allclose(output_adapter_1, output_adapter_2, atol=1e-3, rtol=1e-3), - "Adapter 1 and 2 should give different results", - ) + self.assertFalse( + np.allclose(output_adapter_2, output_adapter_mixed, atol=1e-3, rtol=1e-3), + "Adapter 2 and mixed adapters should give different results", + ) - self.assertFalse( - np.allclose(output_adapter_1, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter 1 and mixed adapters should give different results", - ) + pipe.disable_lora() + output_disabled = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse( - np.allclose(output_adapter_2, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter 2 and mixed adapters should give different results", - ) - - pipe.disable_lora() - output_disabled = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(output_no_lora, output_disabled, atol=1e-3, rtol=1e-3), - "output with no lora and output with lora disabled should give same results", - ) + self.assertTrue( + np.allclose(output_no_lora, output_disabled, atol=1e-3, rtol=1e-3), + "output with no lora and output with lora disabled should give same results", + ) def test_wrong_adapter_name_raises_error(self): adapter_name = "adapter-1" - scheduler_cls = self.scheduler_classes[0] - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() pipe = self.pipeline_class(**components) pipe = pipe.to(torch_device) pipe.set_progress_bar_config(disable=None) @@ -1024,8 +989,7 @@ class PeftLoraLoaderMixinTests: def test_multiple_wrong_adapter_name_raises_error(self): adapter_name = "adapter-1" - scheduler_cls = self.scheduler_classes[0] - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() pipe = self.pipeline_class(**components) pipe = pipe.to(torch_device) pipe.set_progress_bar_config(disable=None) @@ -1054,131 +1018,127 @@ class PeftLoraLoaderMixinTests: Tests a simple inference with lora attached to text encoder and unet, attaches one adapter and set different weights for different blocks (i.e. block lora) """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") + pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") + self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config) - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") + denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet + denoiser.add_adapter(denoiser_lora_config) + self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) + if self.has_two_text_encoders or self.has_three_text_encoders: + if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: + pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") + self.assertTrue( + check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" + ) - weights_1 = {"text_encoder": 2, "unet": {"down": 5}} - pipe.set_adapters("adapter-1", weights_1) - output_weights_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] + weights_1 = {"text_encoder": 2, "unet": {"down": 5}} + pipe.set_adapters("adapter-1", weights_1) + output_weights_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] - weights_2 = {"unet": {"up": 5}} - pipe.set_adapters("adapter-1", weights_2) - output_weights_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] + weights_2 = {"unet": {"up": 5}} + pipe.set_adapters("adapter-1", weights_2) + output_weights_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse( - np.allclose(output_weights_1, output_weights_2, atol=1e-3, rtol=1e-3), - "LoRA weights 1 and 2 should give different results", - ) - self.assertFalse( - np.allclose(output_no_lora, output_weights_1, atol=1e-3, rtol=1e-3), - "No adapter and LoRA weights 1 should give different results", - ) - self.assertFalse( - np.allclose(output_no_lora, output_weights_2, atol=1e-3, rtol=1e-3), - "No adapter and LoRA weights 2 should give different results", - ) + self.assertFalse( + np.allclose(output_weights_1, output_weights_2, atol=1e-3, rtol=1e-3), + "LoRA weights 1 and 2 should give different results", + ) + self.assertFalse( + np.allclose(output_no_lora, output_weights_1, atol=1e-3, rtol=1e-3), + "No adapter and LoRA weights 1 should give different results", + ) + self.assertFalse( + np.allclose(output_no_lora, output_weights_2, atol=1e-3, rtol=1e-3), + "No adapter and LoRA weights 2 should give different results", + ) - pipe.disable_lora() - output_disabled = pipe(**inputs, generator=torch.manual_seed(0))[0] + pipe.disable_lora() + output_disabled = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - np.allclose(output_no_lora, output_disabled, atol=1e-3, rtol=1e-3), - "output with no lora and output with lora disabled should give same results", - ) + self.assertTrue( + np.allclose(output_no_lora, output_disabled, atol=1e-3, rtol=1e-3), + "output with no lora and output with lora disabled should give same results", + ) def test_simple_inference_with_text_denoiser_multi_adapter_block_lora(self): """ Tests a simple inference with lora attached to text encoder and unet, attaches multiple adapters and set different weights for different blocks (i.e. block lora) """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") + if "text_encoder" in self.pipeline_class._lora_loadable_modules: + pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") + pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") + self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") + + denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet + denoiser.add_adapter(denoiser_lora_config, "adapter-1") + denoiser.add_adapter(denoiser_lora_config, "adapter-2") + self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") + + if self.has_two_text_encoders or self.has_three_text_encoders: + if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: + pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") + pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-2") self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder" + check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" ) - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") - denoiser.add_adapter(denoiser_lora_config, "adapter-2") - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") + scales_1 = {"text_encoder": 2, "unet": {"down": 5}} + scales_2 = {"unet": {"down": 5, "mid": 5}} - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-2") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) + pipe.set_adapters("adapter-1", scales_1) + output_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] - scales_1 = {"text_encoder": 2, "unet": {"down": 5}} - scales_2 = {"unet": {"down": 5, "mid": 5}} + pipe.set_adapters("adapter-2", scales_2) + output_adapter_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] - pipe.set_adapters("adapter-1", scales_1) - output_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] + pipe.set_adapters(["adapter-1", "adapter-2"], [scales_1, scales_2]) + output_adapter_mixed = pipe(**inputs, generator=torch.manual_seed(0))[0] - pipe.set_adapters("adapter-2", scales_2) - output_adapter_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] + # Fuse and unfuse should lead to the same results + self.assertFalse( + np.allclose(output_adapter_1, output_adapter_2, atol=1e-3, rtol=1e-3), + "Adapter 1 and 2 should give different results", + ) - pipe.set_adapters(["adapter-1", "adapter-2"], [scales_1, scales_2]) - output_adapter_mixed = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertFalse( + np.allclose(output_adapter_1, output_adapter_mixed, atol=1e-3, rtol=1e-3), + "Adapter 1 and mixed adapters should give different results", + ) - # Fuse and unfuse should lead to the same results - self.assertFalse( - np.allclose(output_adapter_1, output_adapter_2, atol=1e-3, rtol=1e-3), - "Adapter 1 and 2 should give different results", - ) + self.assertFalse( + np.allclose(output_adapter_2, output_adapter_mixed, atol=1e-3, rtol=1e-3), + "Adapter 2 and mixed adapters should give different results", + ) - self.assertFalse( - np.allclose(output_adapter_1, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter 1 and mixed adapters should give different results", - ) + pipe.disable_lora() + output_disabled = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse( - np.allclose(output_adapter_2, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter 2 and mixed adapters should give different results", - ) + self.assertTrue( + np.allclose(output_no_lora, output_disabled, atol=1e-3, rtol=1e-3), + "output with no lora and output with lora disabled should give same results", + ) - pipe.disable_lora() - output_disabled = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(output_no_lora, output_disabled, atol=1e-3, rtol=1e-3), - "output with no lora and output with lora disabled should give same results", - ) - - # a mismatching number of adapter_names and adapter_weights should raise an error - with self.assertRaises(ValueError): - pipe.set_adapters(["adapter-1", "adapter-2"], [scales_1]) + # a mismatching number of adapter_names and adapter_weights should raise an error + with self.assertRaises(ValueError): + pipe.set_adapters(["adapter-1", "adapter-2"], [scales_1]) def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): """Tests that any valid combination of lora block scales can be used in pipe.set_adapter""" @@ -1274,170 +1234,164 @@ class PeftLoraLoaderMixinTests: Tests a simple inference with lora attached to text encoder and unet, attaches multiple adapters and set/delete them """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") + if "text_encoder" in self.pipeline_class._lora_loadable_modules: + pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") + pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") + self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") + + denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet + denoiser.add_adapter(denoiser_lora_config, "adapter-1") + denoiser.add_adapter(denoiser_lora_config, "adapter-2") + self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") + + if self.has_two_text_encoders or self.has_three_text_encoders: + lora_loadable_components = self.pipeline_class._lora_loadable_modules + if "text_encoder_2" in lora_loadable_components: + pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") + pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-2") self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder" + check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" ) - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") - denoiser.add_adapter(denoiser_lora_config, "adapter-2") - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") + pipe.set_adapters("adapter-1") + output_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] - if self.has_two_text_encoders or self.has_three_text_encoders: - lora_loadable_components = self.pipeline_class._lora_loadable_modules - if "text_encoder_2" in lora_loadable_components: - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-2") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) + pipe.set_adapters("adapter-2") + output_adapter_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] - pipe.set_adapters("adapter-1") - output_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] + pipe.set_adapters(["adapter-1", "adapter-2"]) + output_adapter_mixed = pipe(**inputs, generator=torch.manual_seed(0))[0] - pipe.set_adapters("adapter-2") - output_adapter_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertFalse( + np.allclose(output_adapter_1, output_adapter_2, atol=1e-3, rtol=1e-3), + "Adapter 1 and 2 should give different results", + ) - pipe.set_adapters(["adapter-1", "adapter-2"]) - output_adapter_mixed = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertFalse( + np.allclose(output_adapter_1, output_adapter_mixed, atol=1e-3, rtol=1e-3), + "Adapter 1 and mixed adapters should give different results", + ) - self.assertFalse( - np.allclose(output_adapter_1, output_adapter_2, atol=1e-3, rtol=1e-3), - "Adapter 1 and 2 should give different results", - ) + self.assertFalse( + np.allclose(output_adapter_2, output_adapter_mixed, atol=1e-3, rtol=1e-3), + "Adapter 2 and mixed adapters should give different results", + ) - self.assertFalse( - np.allclose(output_adapter_1, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter 1 and mixed adapters should give different results", - ) + pipe.delete_adapters("adapter-1") + output_deleted_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse( - np.allclose(output_adapter_2, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter 2 and mixed adapters should give different results", - ) + self.assertTrue( + np.allclose(output_deleted_adapter_1, output_adapter_2, atol=1e-3, rtol=1e-3), + "Adapter 1 and 2 should give different results", + ) - pipe.delete_adapters("adapter-1") - output_deleted_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] + pipe.delete_adapters("adapter-2") + output_deleted_adapters = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - np.allclose(output_deleted_adapter_1, output_adapter_2, atol=1e-3, rtol=1e-3), - "Adapter 1 and 2 should give different results", - ) + self.assertTrue( + np.allclose(output_no_lora, output_deleted_adapters, atol=1e-3, rtol=1e-3), + "output with no lora and output with lora disabled should give same results", + ) - pipe.delete_adapters("adapter-2") - output_deleted_adapters = pipe(**inputs, generator=torch.manual_seed(0))[0] + if "text_encoder" in self.pipeline_class._lora_loadable_modules: + pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") + pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") - self.assertTrue( - np.allclose(output_no_lora, output_deleted_adapters, atol=1e-3, rtol=1e-3), - "output with no lora and output with lora disabled should give same results", - ) + denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet + denoiser.add_adapter(denoiser_lora_config, "adapter-1") + denoiser.add_adapter(denoiser_lora_config, "adapter-2") + self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") + pipe.set_adapters(["adapter-1", "adapter-2"]) + pipe.delete_adapters(["adapter-1", "adapter-2"]) - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") - denoiser.add_adapter(denoiser_lora_config, "adapter-2") - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") + output_deleted_adapters = pipe(**inputs, generator=torch.manual_seed(0))[0] - pipe.set_adapters(["adapter-1", "adapter-2"]) - pipe.delete_adapters(["adapter-1", "adapter-2"]) - - output_deleted_adapters = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(output_no_lora, output_deleted_adapters, atol=1e-3, rtol=1e-3), - "output with no lora and output with lora disabled should give same results", - ) + self.assertTrue( + np.allclose(output_no_lora, output_deleted_adapters, atol=1e-3, rtol=1e-3), + "output with no lora and output with lora disabled should give same results", + ) def test_simple_inference_with_text_denoiser_multi_adapter_weighted(self): """ Tests a simple inference with lora attached to text encoder and unet, attaches multiple adapters and set them """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") + if "text_encoder" in self.pipeline_class._lora_loadable_modules: + pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") + pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") + self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") + + denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet + denoiser.add_adapter(denoiser_lora_config, "adapter-1") + denoiser.add_adapter(denoiser_lora_config, "adapter-2") + self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") + + if self.has_two_text_encoders or self.has_three_text_encoders: + lora_loadable_components = self.pipeline_class._lora_loadable_modules + if "text_encoder_2" in lora_loadable_components: + pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") + pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-2") self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder" + check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" ) - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") - denoiser.add_adapter(denoiser_lora_config, "adapter-2") - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") + pipe.set_adapters("adapter-1") + output_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] - if self.has_two_text_encoders or self.has_three_text_encoders: - lora_loadable_components = self.pipeline_class._lora_loadable_modules - if "text_encoder_2" in lora_loadable_components: - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-2") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) + pipe.set_adapters("adapter-2") + output_adapter_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] - pipe.set_adapters("adapter-1") - output_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] + pipe.set_adapters(["adapter-1", "adapter-2"]) + output_adapter_mixed = pipe(**inputs, generator=torch.manual_seed(0))[0] - pipe.set_adapters("adapter-2") - output_adapter_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] + # Fuse and unfuse should lead to the same results + self.assertFalse( + np.allclose(output_adapter_1, output_adapter_2, atol=1e-3, rtol=1e-3), + "Adapter 1 and 2 should give different results", + ) - pipe.set_adapters(["adapter-1", "adapter-2"]) - output_adapter_mixed = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertFalse( + np.allclose(output_adapter_1, output_adapter_mixed, atol=1e-3, rtol=1e-3), + "Adapter 1 and mixed adapters should give different results", + ) - # Fuse and unfuse should lead to the same results - self.assertFalse( - np.allclose(output_adapter_1, output_adapter_2, atol=1e-3, rtol=1e-3), - "Adapter 1 and 2 should give different results", - ) + self.assertFalse( + np.allclose(output_adapter_2, output_adapter_mixed, atol=1e-3, rtol=1e-3), + "Adapter 2 and mixed adapters should give different results", + ) - self.assertFalse( - np.allclose(output_adapter_1, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter 1 and mixed adapters should give different results", - ) + pipe.set_adapters(["adapter-1", "adapter-2"], [0.5, 0.6]) + output_adapter_mixed_weighted = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse( - np.allclose(output_adapter_2, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter 2 and mixed adapters should give different results", - ) + self.assertFalse( + np.allclose(output_adapter_mixed_weighted, output_adapter_mixed, atol=1e-3, rtol=1e-3), + "Weighted adapter and mixed adapter should give different results", + ) - pipe.set_adapters(["adapter-1", "adapter-2"], [0.5, 0.6]) - output_adapter_mixed_weighted = pipe(**inputs, generator=torch.manual_seed(0))[0] + pipe.disable_lora() + output_disabled = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse( - np.allclose(output_adapter_mixed_weighted, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Weighted adapter and mixed adapter should give different results", - ) - - pipe.disable_lora() - output_disabled = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(output_no_lora, output_disabled, atol=1e-3, rtol=1e-3), - "output with no lora and output with lora disabled should give same results", - ) + self.assertTrue( + np.allclose(output_no_lora, output_disabled, atol=1e-3, rtol=1e-3), + "output with no lora and output with lora disabled should give same results", + ) @skip_mps @pytest.mark.xfail( @@ -1446,163 +1400,156 @@ class PeftLoraLoaderMixinTests: strict=False, ) def test_lora_fuse_nan(self): - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder" + if "text_encoder" in self.pipeline_class._lora_loadable_modules: + pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") + self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") + + denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet + denoiser.add_adapter(denoiser_lora_config, "adapter-1") + self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") + + # corrupt one LoRA weight with `inf` values + with torch.no_grad(): + if self.unet_kwargs: + pipe.unet.mid_block.attentions[0].transformer_blocks[0].attn1.to_q.lora_A["adapter-1"].weight += float( + "inf" ) + else: + named_modules = [name for name, _ in pipe.transformer.named_modules()] + possible_tower_names = [ + "transformer_blocks", + "blocks", + "joint_transformer_blocks", + "single_transformer_blocks", + ] + filtered_tower_names = [ + tower_name for tower_name in possible_tower_names if hasattr(pipe.transformer, tower_name) + ] + if len(filtered_tower_names) == 0: + reason = f"`pipe.transformer` didn't have any of the following attributes: {possible_tower_names}." + raise ValueError(reason) + for tower_name in filtered_tower_names: + transformer_tower = getattr(pipe.transformer, tower_name) + has_attn1 = any("attn1" in name for name in named_modules) + if has_attn1: + transformer_tower[0].attn1.to_q.lora_A["adapter-1"].weight += float("inf") + else: + transformer_tower[0].attn.to_q.lora_A["adapter-1"].weight += float("inf") - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") + # with `safe_fusing=True` we should see an Error + with self.assertRaises(ValueError): + pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules, safe_fusing=True) - # corrupt one LoRA weight with `inf` values - with torch.no_grad(): - if self.unet_kwargs: - pipe.unet.mid_block.attentions[0].transformer_blocks[0].attn1.to_q.lora_A[ - "adapter-1" - ].weight += float("inf") - else: - named_modules = [name for name, _ in pipe.transformer.named_modules()] - possible_tower_names = [ - "transformer_blocks", - "blocks", - "joint_transformer_blocks", - "single_transformer_blocks", - ] - filtered_tower_names = [ - tower_name for tower_name in possible_tower_names if hasattr(pipe.transformer, tower_name) - ] - if len(filtered_tower_names) == 0: - reason = ( - f"`pipe.transformer` didn't have any of the following attributes: {possible_tower_names}." - ) - raise ValueError(reason) - for tower_name in filtered_tower_names: - transformer_tower = getattr(pipe.transformer, tower_name) - has_attn1 = any("attn1" in name for name in named_modules) - if has_attn1: - transformer_tower[0].attn1.to_q.lora_A["adapter-1"].weight += float("inf") - else: - transformer_tower[0].attn.to_q.lora_A["adapter-1"].weight += float("inf") + # without we should not see an error, but every image will be black + pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules, safe_fusing=False) + out = pipe(**inputs)[0] - # with `safe_fusing=True` we should see an Error - with self.assertRaises(ValueError): - pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules, safe_fusing=True) - - # without we should not see an error, but every image will be black - pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules, safe_fusing=False) - out = pipe(**inputs)[0] - - self.assertTrue(np.isnan(out).all()) + self.assertTrue(np.isnan(out).all()) def test_get_adapters(self): """ Tests a simple usecase where we attach multiple adapters and check if the results are the expected results """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") + pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") + denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet + denoiser.add_adapter(denoiser_lora_config, "adapter-1") - adapter_names = pipe.get_active_adapters() - self.assertListEqual(adapter_names, ["adapter-1"]) + adapter_names = pipe.get_active_adapters() + self.assertListEqual(adapter_names, ["adapter-1"]) - pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") - denoiser.add_adapter(denoiser_lora_config, "adapter-2") + pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") + denoiser.add_adapter(denoiser_lora_config, "adapter-2") - adapter_names = pipe.get_active_adapters() - self.assertListEqual(adapter_names, ["adapter-2"]) + adapter_names = pipe.get_active_adapters() + self.assertListEqual(adapter_names, ["adapter-2"]) - pipe.set_adapters(["adapter-1", "adapter-2"]) - self.assertListEqual(pipe.get_active_adapters(), ["adapter-1", "adapter-2"]) + pipe.set_adapters(["adapter-1", "adapter-2"]) + self.assertListEqual(pipe.get_active_adapters(), ["adapter-1", "adapter-2"]) def test_get_list_adapters(self): """ Tests a simple usecase where we attach multiple adapters and check if the results are the expected results """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) - # 1. - dicts_to_be_checked = {} - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - dicts_to_be_checked = {"text_encoder": ["adapter-1"]} + # 1. + dicts_to_be_checked = {} + if "text_encoder" in self.pipeline_class._lora_loadable_modules: + pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") + dicts_to_be_checked = {"text_encoder": ["adapter-1"]} - if self.unet_kwargs is not None: - pipe.unet.add_adapter(denoiser_lora_config, "adapter-1") - dicts_to_be_checked.update({"unet": ["adapter-1"]}) - else: - pipe.transformer.add_adapter(denoiser_lora_config, "adapter-1") - dicts_to_be_checked.update({"transformer": ["adapter-1"]}) + if self.unet_kwargs is not None: + pipe.unet.add_adapter(denoiser_lora_config, "adapter-1") + dicts_to_be_checked.update({"unet": ["adapter-1"]}) + else: + pipe.transformer.add_adapter(denoiser_lora_config, "adapter-1") + dicts_to_be_checked.update({"transformer": ["adapter-1"]}) - self.assertDictEqual(pipe.get_list_adapters(), dicts_to_be_checked) + self.assertDictEqual(pipe.get_list_adapters(), dicts_to_be_checked) - # 2. - dicts_to_be_checked = {} - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") - dicts_to_be_checked = {"text_encoder": ["adapter-1", "adapter-2"]} + # 2. + dicts_to_be_checked = {} + if "text_encoder" in self.pipeline_class._lora_loadable_modules: + pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") + dicts_to_be_checked = {"text_encoder": ["adapter-1", "adapter-2"]} - if self.unet_kwargs is not None: - pipe.unet.add_adapter(denoiser_lora_config, "adapter-2") - dicts_to_be_checked.update({"unet": ["adapter-1", "adapter-2"]}) - else: - pipe.transformer.add_adapter(denoiser_lora_config, "adapter-2") - dicts_to_be_checked.update({"transformer": ["adapter-1", "adapter-2"]}) + if self.unet_kwargs is not None: + pipe.unet.add_adapter(denoiser_lora_config, "adapter-2") + dicts_to_be_checked.update({"unet": ["adapter-1", "adapter-2"]}) + else: + pipe.transformer.add_adapter(denoiser_lora_config, "adapter-2") + dicts_to_be_checked.update({"transformer": ["adapter-1", "adapter-2"]}) - self.assertDictEqual(pipe.get_list_adapters(), dicts_to_be_checked) + self.assertDictEqual(pipe.get_list_adapters(), dicts_to_be_checked) - # 3. - pipe.set_adapters(["adapter-1", "adapter-2"]) + # 3. + pipe.set_adapters(["adapter-1", "adapter-2"]) - dicts_to_be_checked = {} - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - dicts_to_be_checked = {"text_encoder": ["adapter-1", "adapter-2"]} + dicts_to_be_checked = {} + if "text_encoder" in self.pipeline_class._lora_loadable_modules: + dicts_to_be_checked = {"text_encoder": ["adapter-1", "adapter-2"]} - if self.unet_kwargs is not None: - dicts_to_be_checked.update({"unet": ["adapter-1", "adapter-2"]}) - else: - dicts_to_be_checked.update({"transformer": ["adapter-1", "adapter-2"]}) + if self.unet_kwargs is not None: + dicts_to_be_checked.update({"unet": ["adapter-1", "adapter-2"]}) + else: + dicts_to_be_checked.update({"transformer": ["adapter-1", "adapter-2"]}) - self.assertDictEqual( - pipe.get_list_adapters(), - dicts_to_be_checked, - ) + self.assertDictEqual( + pipe.get_list_adapters(), + dicts_to_be_checked, + ) - # 4. - dicts_to_be_checked = {} - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - dicts_to_be_checked = {"text_encoder": ["adapter-1", "adapter-2"]} + # 4. + dicts_to_be_checked = {} + if "text_encoder" in self.pipeline_class._lora_loadable_modules: + dicts_to_be_checked = {"text_encoder": ["adapter-1", "adapter-2"]} - if self.unet_kwargs is not None: - pipe.unet.add_adapter(denoiser_lora_config, "adapter-3") - dicts_to_be_checked.update({"unet": ["adapter-1", "adapter-2", "adapter-3"]}) - else: - pipe.transformer.add_adapter(denoiser_lora_config, "adapter-3") - dicts_to_be_checked.update({"transformer": ["adapter-1", "adapter-2", "adapter-3"]}) + if self.unet_kwargs is not None: + pipe.unet.add_adapter(denoiser_lora_config, "adapter-3") + dicts_to_be_checked.update({"unet": ["adapter-1", "adapter-2", "adapter-3"]}) + else: + pipe.transformer.add_adapter(denoiser_lora_config, "adapter-3") + dicts_to_be_checked.update({"transformer": ["adapter-1", "adapter-2", "adapter-3"]}) - self.assertDictEqual(pipe.get_list_adapters(), dicts_to_be_checked) + self.assertDictEqual(pipe.get_list_adapters(), dicts_to_be_checked) @require_peft_version_greater(peft_version="0.6.2") def test_simple_inference_with_text_lora_denoiser_fused_multi( @@ -1612,8 +1559,83 @@ class PeftLoraLoaderMixinTests: Tests a simple inference with lora attached into text encoder + fuses the lora weights into base model and makes sure it works as expected - with unet and multi-adapter case """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) + + output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue(output_no_lora.shape == self.output_shape) + + if "text_encoder" in self.pipeline_class._lora_loadable_modules: + pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") + self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") + pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") + + denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet + denoiser.add_adapter(denoiser_lora_config, "adapter-1") + self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") + denoiser.add_adapter(denoiser_lora_config, "adapter-2") + + if self.has_two_text_encoders or self.has_three_text_encoders: + lora_loadable_components = self.pipeline_class._lora_loadable_modules + if "text_encoder_2" in lora_loadable_components: + pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") + self.assertTrue( + check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" + ) + pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-2") + + # set them to multi-adapter inference mode + pipe.set_adapters(["adapter-1", "adapter-2"]) + outputs_all_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + + pipe.set_adapters(["adapter-1"]) + outputs_lora_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] + + pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules, adapter_names=["adapter-1"]) + self.assertTrue(pipe.num_fused_loras == 1, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") + + # Fusing should still keep the LoRA layers so output should remain the same + outputs_lora_1_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] + + self.assertTrue( + np.allclose(outputs_lora_1, outputs_lora_1_fused, atol=expected_atol, rtol=expected_rtol), + "Fused lora should not change the output", + ) + + pipe.unfuse_lora(components=self.pipeline_class._lora_loadable_modules) + self.assertTrue(pipe.num_fused_loras == 0, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") + + if "text_encoder" in self.pipeline_class._lora_loadable_modules: + self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Unfuse should still keep LoRA layers") + + self.assertTrue(check_if_lora_correctly_set(denoiser), "Unfuse should still keep LoRA layers") + + if self.has_two_text_encoders or self.has_three_text_encoders: + if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: + self.assertTrue( + check_if_lora_correctly_set(pipe.text_encoder_2), "Unfuse should still keep LoRA layers" + ) + + pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules, adapter_names=["adapter-2", "adapter-1"]) + self.assertTrue(pipe.num_fused_loras == 2, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") + + # Fusing should still keep the LoRA layers + output_all_lora_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue( + np.allclose(output_all_lora_fused, outputs_all_lora, atol=expected_atol, rtol=expected_rtol), + "Fused lora should not change the output", + ) + pipe.unfuse_lora(components=self.pipeline_class._lora_loadable_modules) + self.assertTrue(pipe.num_fused_loras == 0, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") + + def test_lora_scale_kwargs_match_fusion(self, expected_atol: float = 1e-3, expected_rtol: float = 1e-3): + attention_kwargs_name = determine_attention_kwargs_name(self.pipeline_class) + + for lora_scale in [1.0, 0.8]: + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() pipe = self.pipeline_class(**components) pipe = pipe.to(torch_device) pipe.set_progress_bar_config(disable=None) @@ -1627,150 +1649,65 @@ class PeftLoraLoaderMixinTests: self.assertTrue( check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder" ) - pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet denoiser.add_adapter(denoiser_lora_config, "adapter-1") self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - denoiser.add_adapter(denoiser_lora_config, "adapter-2") if self.has_two_text_encoders or self.has_three_text_encoders: lora_loadable_components = self.pipeline_class._lora_loadable_modules if "text_encoder_2" in lora_loadable_components: pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" + check_if_lora_correctly_set(pipe.text_encoder_2), + "Lora not correctly set in text encoder 2", ) - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-2") - - # set them to multi-adapter inference mode - pipe.set_adapters(["adapter-1", "adapter-2"]) - outputs_all_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] pipe.set_adapters(["adapter-1"]) - outputs_lora_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] + attention_kwargs = {attention_kwargs_name: {"scale": lora_scale}} + outputs_lora_1 = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] - pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules, adapter_names=["adapter-1"]) + pipe.fuse_lora( + components=self.pipeline_class._lora_loadable_modules, + adapter_names=["adapter-1"], + lora_scale=lora_scale, + ) self.assertTrue(pipe.num_fused_loras == 1, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") - # Fusing should still keep the LoRA layers so output should remain the same outputs_lora_1_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] self.assertTrue( np.allclose(outputs_lora_1, outputs_lora_1_fused, atol=expected_atol, rtol=expected_rtol), "Fused lora should not change the output", ) - - pipe.unfuse_lora(components=self.pipeline_class._lora_loadable_modules) - self.assertTrue(pipe.num_fused_loras == 0, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") - - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Unfuse should still keep LoRA layers") - - self.assertTrue(check_if_lora_correctly_set(denoiser), "Unfuse should still keep LoRA layers") - - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Unfuse should still keep LoRA layers" - ) - - pipe.fuse_lora( - components=self.pipeline_class._lora_loadable_modules, adapter_names=["adapter-2", "adapter-1"] + self.assertFalse( + np.allclose(output_no_lora, outputs_lora_1, atol=expected_atol, rtol=expected_rtol), + "LoRA should change the output", ) - self.assertTrue(pipe.num_fused_loras == 2, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") - - # Fusing should still keep the LoRA layers - output_all_lora_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - np.allclose(output_all_lora_fused, outputs_all_lora, atol=expected_atol, rtol=expected_rtol), - "Fused lora should not change the output", - ) - pipe.unfuse_lora(components=self.pipeline_class._lora_loadable_modules) - self.assertTrue(pipe.num_fused_loras == 0, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") - - def test_lora_scale_kwargs_match_fusion(self, expected_atol: float = 1e-3, expected_rtol: float = 1e-3): - attention_kwargs_name = determine_attention_kwargs_name(self.pipeline_class) - - for lora_scale in [1.0, 0.8]: - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) - - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder" - ) - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - - if self.has_two_text_encoders or self.has_three_text_encoders: - lora_loadable_components = self.pipeline_class._lora_loadable_modules - if "text_encoder_2" in lora_loadable_components: - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), - "Lora not correctly set in text encoder 2", - ) - - pipe.set_adapters(["adapter-1"]) - attention_kwargs = {attention_kwargs_name: {"scale": lora_scale}} - outputs_lora_1 = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] - - pipe.fuse_lora( - components=self.pipeline_class._lora_loadable_modules, - adapter_names=["adapter-1"], - lora_scale=lora_scale, - ) - self.assertTrue(pipe.num_fused_loras == 1, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") - - outputs_lora_1_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(outputs_lora_1, outputs_lora_1_fused, atol=expected_atol, rtol=expected_rtol), - "Fused lora should not change the output", - ) - self.assertFalse( - np.allclose(output_no_lora, outputs_lora_1, atol=expected_atol, rtol=expected_rtol), - "LoRA should change the output", - ) @require_peft_version_greater(peft_version="0.9.0") def test_simple_inference_with_dora(self): - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components( - scheduler_cls, use_dora=True - ) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components(use_dora=True) + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_dora_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_dora_lora.shape == self.output_shape) + output_no_dora_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue(output_no_dora_lora.shape == self.output_shape) - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) + pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - output_dora_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + output_dora_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse( - np.allclose(output_dora_lora, output_no_dora_lora, atol=1e-3, rtol=1e-3), - "DoRA lora should change the output", - ) + self.assertFalse( + np.allclose(output_dora_lora, output_no_dora_lora, atol=1e-3, rtol=1e-3), + "DoRA lora should change the output", + ) def test_missing_keys_warning(self): - scheduler_cls = self.scheduler_classes[0] # Skip text encoder check for now as that is handled with `transformers`. - components, _, denoiser_lora_config = self.get_dummy_components(scheduler_cls) + components, _, denoiser_lora_config = self.get_dummy_components() pipe = self.pipeline_class(**components) pipe = pipe.to(torch_device) pipe.set_progress_bar_config(disable=None) @@ -1805,9 +1742,8 @@ class PeftLoraLoaderMixinTests: self.assertTrue(missing_key.replace(f"{component}.", "") in cap_logger.out.replace("default_0.", "")) def test_unexpected_keys_warning(self): - scheduler_cls = self.scheduler_classes[0] # Skip text encoder check for now as that is handled with `transformers`. - components, _, denoiser_lora_config = self.get_dummy_components(scheduler_cls) + components, _, denoiser_lora_config = self.get_dummy_components() pipe = self.pipeline_class(**components) pipe = pipe.to(torch_device) pipe.set_progress_bar_config(disable=None) @@ -1842,23 +1778,22 @@ class PeftLoraLoaderMixinTests: Tests a simple inference with lora attached to text encoder and unet, then unloads the lora weights and makes sure it works as expected """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) + pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True) - pipe.text_encoder = torch.compile(pipe.text_encoder, mode="reduce-overhead", fullgraph=True) + pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True) + pipe.text_encoder = torch.compile(pipe.text_encoder, mode="reduce-overhead", fullgraph=True) - if self.has_two_text_encoders or self.has_three_text_encoders: - pipe.text_encoder_2 = torch.compile(pipe.text_encoder_2, mode="reduce-overhead", fullgraph=True) + if self.has_two_text_encoders or self.has_three_text_encoders: + pipe.text_encoder_2 = torch.compile(pipe.text_encoder_2, mode="reduce-overhead", fullgraph=True) - # Just makes sure it works.. - _ = pipe(**inputs, generator=torch.manual_seed(0))[0] + # Just makes sure it works. + _ = pipe(**inputs, generator=torch.manual_seed(0))[0] def test_modify_padding_mode(self): def set_pad_mode(network, mode="circular"): @@ -1866,22 +1801,20 @@ class PeftLoraLoaderMixinTests: if isinstance(module, torch.nn.Conv2d): module.padding_mode = mode - for scheduler_cls in self.scheduler_classes: - components, _, _ = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _pad_mode = "circular" - set_pad_mode(pipe.vae, _pad_mode) - set_pad_mode(pipe.unet, _pad_mode) + components, _, _ = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _pad_mode = "circular" + set_pad_mode(pipe.vae, _pad_mode) + set_pad_mode(pipe.unet, _pad_mode) - _, _, inputs = self.get_dummy_inputs() - _ = pipe(**inputs)[0] + _, _, inputs = self.get_dummy_inputs() + _ = pipe(**inputs)[0] def test_logs_info_when_no_lora_keys_found(self): - scheduler_cls = self.scheduler_classes[0] # Skip text encoder check for now as that is handled with `transformers`. - components, _, _ = self.get_dummy_components(scheduler_cls) + components, _, _ = self.get_dummy_components() pipe = self.pipeline_class(**components) pipe = pipe.to(torch_device) pipe.set_progress_bar_config(disable=None) @@ -1925,73 +1858,71 @@ class PeftLoraLoaderMixinTests: def test_set_adapters_match_attention_kwargs(self): """Test to check if outputs after `set_adapters()` and attention kwargs match.""" attention_kwargs_name = determine_attention_kwargs_name(self.pipeline_class) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) + output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue(output_no_lora.shape == self.output_shape) + + pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) + + lora_scale = 0.5 + attention_kwargs = {attention_kwargs_name: {"scale": lora_scale}} + output_lora_scale = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] + self.assertFalse( + np.allclose(output_no_lora, output_lora_scale, atol=1e-3, rtol=1e-3), + "Lora + scale should change the output", + ) + + pipe.set_adapters("default", lora_scale) + output_lora_scale_wo_kwargs = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue( + not np.allclose(output_no_lora, output_lora_scale_wo_kwargs, atol=1e-3, rtol=1e-3), + "Lora + scale should change the output", + ) + self.assertTrue( + np.allclose(output_lora_scale, output_lora_scale_wo_kwargs, atol=1e-3, rtol=1e-3), + "Lora + scale should match the output of `set_adapters()`.", + ) + + with tempfile.TemporaryDirectory() as tmpdirname: + modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True) + lora_state_dicts = self._get_lora_state_dicts(modules_to_save) + self.pipeline_class.save_lora_weights( + save_directory=tmpdirname, safe_serialization=True, **lora_state_dicts + ) + + self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors"))) pipe = self.pipeline_class(**components) pipe = pipe.to(torch_device) pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors")) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + for module_name, module in modules_to_save.items(): + self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - - lora_scale = 0.5 - attention_kwargs = {attention_kwargs_name: {"scale": lora_scale}} - output_lora_scale = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] - self.assertFalse( - np.allclose(output_no_lora, output_lora_scale, atol=1e-3, rtol=1e-3), - "Lora + scale should change the output", - ) - - pipe.set_adapters("default", lora_scale) - output_lora_scale_wo_kwargs = pipe(**inputs, generator=torch.manual_seed(0))[0] + output_lora_from_pretrained = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] self.assertTrue( - not np.allclose(output_no_lora, output_lora_scale_wo_kwargs, atol=1e-3, rtol=1e-3), + not np.allclose(output_no_lora, output_lora_from_pretrained, atol=1e-3, rtol=1e-3), "Lora + scale should change the output", ) self.assertTrue( - np.allclose(output_lora_scale, output_lora_scale_wo_kwargs, atol=1e-3, rtol=1e-3), - "Lora + scale should match the output of `set_adapters()`.", + np.allclose(output_lora_scale, output_lora_from_pretrained, atol=1e-3, rtol=1e-3), + "Loading from saved checkpoints should give same results as attention_kwargs.", + ) + self.assertTrue( + np.allclose(output_lora_scale_wo_kwargs, output_lora_from_pretrained, atol=1e-3, rtol=1e-3), + "Loading from saved checkpoints should give same results as set_adapters().", ) - - with tempfile.TemporaryDirectory() as tmpdirname: - modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True) - lora_state_dicts = self._get_lora_state_dicts(modules_to_save) - self.pipeline_class.save_lora_weights( - save_directory=tmpdirname, safe_serialization=True, **lora_state_dicts - ) - - self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors"))) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors")) - - for module_name, module in modules_to_save.items(): - self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") - - output_lora_from_pretrained = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] - self.assertTrue( - not np.allclose(output_no_lora, output_lora_from_pretrained, atol=1e-3, rtol=1e-3), - "Lora + scale should change the output", - ) - self.assertTrue( - np.allclose(output_lora_scale, output_lora_from_pretrained, atol=1e-3, rtol=1e-3), - "Loading from saved checkpoints should give same results as attention_kwargs.", - ) - self.assertTrue( - np.allclose(output_lora_scale_wo_kwargs, output_lora_from_pretrained, atol=1e-3, rtol=1e-3), - "Loading from saved checkpoints should give same results as set_adapters().", - ) @require_peft_version_greater("0.13.2") def test_lora_B_bias(self): # Currently, this test is only relevant for Flux Control LoRA as we are not # aware of any other LoRA checkpoint that has its `lora_B` biases trained. - components, _, denoiser_lora_config = self.get_dummy_components(self.scheduler_classes[0]) + components, _, denoiser_lora_config = self.get_dummy_components() pipe = self.pipeline_class(**components) pipe = pipe.to(torch_device) pipe.set_progress_bar_config(disable=None) @@ -2028,7 +1959,7 @@ class PeftLoraLoaderMixinTests: self.assertFalse(np.allclose(lora_bias_false_output, lora_bias_true_output, atol=1e-3, rtol=1e-3)) def test_correct_lora_configs_with_different_ranks(self): - components, _, denoiser_lora_config = self.get_dummy_components(self.scheduler_classes[0]) + components, _, denoiser_lora_config = self.get_dummy_components() pipe = self.pipeline_class(**components) pipe = pipe.to(torch_device) pipe.set_progress_bar_config(disable=None) @@ -2114,7 +2045,7 @@ class PeftLoraLoaderMixinTests: self.assertEqual(submodule.bias.dtype, dtype_to_check) def initialize_pipeline(storage_dtype=None, compute_dtype=torch.float32): - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(self.scheduler_classes[0]) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() pipe = self.pipeline_class(**components) pipe = pipe.to(torch_device, dtype=compute_dtype) pipe.set_progress_bar_config(disable=None) @@ -2181,7 +2112,7 @@ class PeftLoraLoaderMixinTests: self.assertTrue(module._diffusers_hook.get_hook(_PEFT_AUTOCAST_DISABLE_HOOK) is not None) # 1. Test forward with add_adapter - components, _, denoiser_lora_config = self.get_dummy_components(self.scheduler_classes[0]) + components, _, denoiser_lora_config = self.get_dummy_components() pipe = self.pipeline_class(**components) pipe = pipe.to(torch_device, dtype=compute_dtype) pipe.set_progress_bar_config(disable=None) @@ -2211,7 +2142,7 @@ class PeftLoraLoaderMixinTests: ) self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors"))) - components, _, _ = self.get_dummy_components(self.scheduler_classes[0]) + components, _, _ = self.get_dummy_components() pipe = self.pipeline_class(**components) pipe = pipe.to(torch_device, dtype=compute_dtype) pipe.set_progress_bar_config(disable=None) @@ -2231,10 +2162,7 @@ class PeftLoraLoaderMixinTests: @parameterized.expand([4, 8, 16]) def test_lora_adapter_metadata_is_loaded_correctly(self, lora_alpha): - scheduler_cls = self.scheduler_classes[0] - components, text_lora_config, denoiser_lora_config = self.get_dummy_components( - scheduler_cls, lora_alpha=lora_alpha - ) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components(lora_alpha=lora_alpha) pipe = self.pipeline_class(**components) pipe, _ = self.add_adapters_to_pipeline( @@ -2280,10 +2208,7 @@ class PeftLoraLoaderMixinTests: @parameterized.expand([4, 8, 16]) def test_lora_adapter_metadata_save_load_inference(self, lora_alpha): - scheduler_cls = self.scheduler_classes[0] - components, text_lora_config, denoiser_lora_config = self.get_dummy_components( - scheduler_cls, lora_alpha=lora_alpha - ) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components(lora_alpha=lora_alpha) pipe = self.pipeline_class(**components).to(torch_device) _, _, inputs = self.get_dummy_inputs(with_generator=False) @@ -2311,8 +2236,7 @@ class PeftLoraLoaderMixinTests: def test_lora_unload_add_adapter(self): """Tests if `unload_lora_weights()` -> `add_adapter()` works.""" - scheduler_cls = self.scheduler_classes[0] - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() pipe = self.pipeline_class(**components).to(torch_device) _, _, inputs = self.get_dummy_inputs(with_generator=False) @@ -2330,51 +2254,48 @@ class PeftLoraLoaderMixinTests: def test_inference_load_delete_load_adapters(self): "Tests if `load_lora_weights()` -> `delete_adapters()` -> `load_lora_weights()` works." - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config) + if "text_encoder" in self.pipeline_class._lora_loadable_modules: + pipe.text_encoder.add_adapter(text_lora_config) + self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") + + denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet + denoiser.add_adapter(denoiser_lora_config) + self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") + + if self.has_two_text_encoders or self.has_three_text_encoders: + lora_loadable_components = self.pipeline_class._lora_loadable_modules + if "text_encoder_2" in lora_loadable_components: + pipe.text_encoder_2.add_adapter(text_lora_config) self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder" + check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" ) - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config) - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") + output_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] - if self.has_two_text_encoders or self.has_three_text_encoders: - lora_loadable_components = self.pipeline_class._lora_loadable_modules - if "text_encoder_2" in lora_loadable_components: - pipe.text_encoder_2.add_adapter(text_lora_config) - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) + with tempfile.TemporaryDirectory() as tmpdirname: + modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True) + lora_state_dicts = self._get_lora_state_dicts(modules_to_save) + self.pipeline_class.save_lora_weights(save_directory=tmpdirname, **lora_state_dicts) + self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors"))) - output_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] + # First, delete adapter and compare. + pipe.delete_adapters(pipe.get_active_adapters()[0]) + output_no_adapter = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertFalse(np.allclose(output_adapter_1, output_no_adapter, atol=1e-3, rtol=1e-3)) + self.assertTrue(np.allclose(output_no_lora, output_no_adapter, atol=1e-3, rtol=1e-3)) - with tempfile.TemporaryDirectory() as tmpdirname: - modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True) - lora_state_dicts = self._get_lora_state_dicts(modules_to_save) - self.pipeline_class.save_lora_weights(save_directory=tmpdirname, **lora_state_dicts) - self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors"))) - - # First, delete adapter and compare. - pipe.delete_adapters(pipe.get_active_adapters()[0]) - output_no_adapter = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse(np.allclose(output_adapter_1, output_no_adapter, atol=1e-3, rtol=1e-3)) - self.assertTrue(np.allclose(output_no_lora, output_no_adapter, atol=1e-3, rtol=1e-3)) - - # Then load adapter and compare. - pipe.load_lora_weights(tmpdirname) - output_lora_loaded = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(np.allclose(output_adapter_1, output_lora_loaded, atol=1e-3, rtol=1e-3)) + # Then load adapter and compare. + pipe.load_lora_weights(tmpdirname) + output_lora_loaded = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue(np.allclose(output_adapter_1, output_lora_loaded, atol=1e-3, rtol=1e-3)) def _test_group_offloading_inference_denoiser(self, offload_type, use_stream): from diffusers.hooks.group_offloading import _get_top_level_group_offload_hook @@ -2382,7 +2303,7 @@ class PeftLoraLoaderMixinTests: onload_device = torch_device offload_device = torch.device("cpu") - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(self.scheduler_classes[0]) + components, text_lora_config, denoiser_lora_config = self.get_dummy_components() pipe = self.pipeline_class(**components) pipe = pipe.to(torch_device) pipe.set_progress_bar_config(disable=None) @@ -2399,7 +2320,7 @@ class PeftLoraLoaderMixinTests: ) self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors"))) - components, _, _ = self.get_dummy_components(self.scheduler_classes[0]) + components, _, _ = self.get_dummy_components() pipe = self.pipeline_class(**components) pipe.set_progress_bar_config(disable=None) denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet @@ -2451,7 +2372,7 @@ class PeftLoraLoaderMixinTests: @require_torch_accelerator def test_lora_loading_model_cpu_offload(self): - components, _, denoiser_lora_config = self.get_dummy_components(self.scheduler_classes[0]) + components, _, denoiser_lora_config = self.get_dummy_components() _, _, inputs = self.get_dummy_inputs(with_generator=False) pipe = self.pipeline_class(**components) pipe = pipe.to(torch_device) @@ -2470,7 +2391,7 @@ class PeftLoraLoaderMixinTests: save_directory=tmpdirname, safe_serialization=True, **lora_state_dicts ) # reinitialize the pipeline to mimic the inference workflow. - components, _, denoiser_lora_config = self.get_dummy_components(self.scheduler_classes[0]) + components, _, denoiser_lora_config = self.get_dummy_components() pipe = self.pipeline_class(**components) pipe.enable_model_cpu_offload(device=torch_device) pipe.load_lora_weights(tmpdirname) From 7c54a7b38a9ce42db7c6d8ec86ec656cac9ee216 Mon Sep 17 00:00:00 2001 From: Dhruv Nair Date: Wed, 24 Sep 2025 05:23:41 +0200 Subject: [PATCH 020/123] Fix Custom Code loading (#12378) * update * update * update --- src/diffusers/models/auto_model.py | 2 +- src/diffusers/modular_pipelines/modular_pipeline.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/diffusers/models/auto_model.py b/src/diffusers/models/auto_model.py index ada0d54e54..47f3a992b3 100644 --- a/src/diffusers/models/auto_model.py +++ b/src/diffusers/models/auto_model.py @@ -194,7 +194,7 @@ class AutoModel(ConfigMixin): has_remote_code = "auto_map" in config and cls.__name__ in config["auto_map"] trust_remote_code = resolve_trust_remote_code(trust_remote_code, pretrained_model_or_path, has_remote_code) - if not (has_remote_code and trust_remote_code): + if not has_remote_code and trust_remote_code: raise ValueError( "Selected model repository does not happear to have any custom code or does not have a valid `config.json` file." ) diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 78226a49b1..74ffc62348 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -323,7 +323,7 @@ class ModularPipelineBlocks(ConfigMixin, PushToHubMixin): trust_remote_code = resolve_trust_remote_code( trust_remote_code, pretrained_model_name_or_path, has_remote_code ) - if not (has_remote_code and trust_remote_code): + if not has_remote_code and trust_remote_code: raise ValueError( "Selected model repository does not happear to have any custom code or does not have a valid `config.json` file." ) From 9ef118509e1a2682bb6870533cc3552726284167 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Wed, 24 Sep 2025 09:02:25 +0530 Subject: [PATCH 021/123] [tests] disable xformer tests for pipelines it isn't popular. (#12277) disable xformer tests for pipelines it isn't popular. --- tests/pipelines/easyanimate/test_easyanimate.py | 1 + tests/pipelines/hidream_image/test_pipeline_hidream.py | 2 +- tests/pipelines/omnigen/test_pipeline_omnigen.py | 2 +- tests/pipelines/qwenimage/test_qwenimage_controlnet.py | 3 +-- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/pipelines/easyanimate/test_easyanimate.py b/tests/pipelines/easyanimate/test_easyanimate.py index 2dbb8639f1..5cb2a232bb 100644 --- a/tests/pipelines/easyanimate/test_easyanimate.py +++ b/tests/pipelines/easyanimate/test_easyanimate.py @@ -48,6 +48,7 @@ class EasyAnimatePipelineFastTests(PipelineTesterMixin, unittest.TestCase): batch_params = TEXT_TO_IMAGE_BATCH_PARAMS image_params = TEXT_TO_IMAGE_IMAGE_PARAMS image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS + test_xformers_attention = False required_optional_params = frozenset( [ "num_inference_steps", diff --git a/tests/pipelines/hidream_image/test_pipeline_hidream.py b/tests/pipelines/hidream_image/test_pipeline_hidream.py index ec8d36e1d3..ddf39ba4c1 100644 --- a/tests/pipelines/hidream_image/test_pipeline_hidream.py +++ b/tests/pipelines/hidream_image/test_pipeline_hidream.py @@ -47,8 +47,8 @@ class HiDreamImagePipelineFastTests(PipelineTesterMixin, unittest.TestCase): batch_params = TEXT_TO_IMAGE_BATCH_PARAMS image_params = TEXT_TO_IMAGE_IMAGE_PARAMS image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS - required_optional_params = PipelineTesterMixin.required_optional_params + test_xformers_attention = False test_layerwise_casting = True supports_dduf = False diff --git a/tests/pipelines/omnigen/test_pipeline_omnigen.py b/tests/pipelines/omnigen/test_pipeline_omnigen.py index 28648aa76f..1a758b7050 100644 --- a/tests/pipelines/omnigen/test_pipeline_omnigen.py +++ b/tests/pipelines/omnigen/test_pipeline_omnigen.py @@ -22,7 +22,7 @@ class OmniGenPipelineFastTests(unittest.TestCase, PipelineTesterMixin): pipeline_class = OmniGenPipeline params = frozenset(["prompt", "guidance_scale"]) batch_params = frozenset(["prompt"]) - + test_xformers_attention = False test_layerwise_casting = True def get_dummy_components(self): diff --git a/tests/pipelines/qwenimage/test_qwenimage_controlnet.py b/tests/pipelines/qwenimage/test_qwenimage_controlnet.py index c78e5cb233..188106b49b 100644 --- a/tests/pipelines/qwenimage/test_qwenimage_controlnet.py +++ b/tests/pipelines/qwenimage/test_qwenimage_controlnet.py @@ -44,7 +44,6 @@ class QwenControlNetPipelineFastTests(PipelineTesterMixin, unittest.TestCase): batch_params = frozenset(["prompt", "negative_prompt", "control_image"]) image_params = frozenset(["control_image"]) image_latents_params = frozenset(["latents"]) - required_optional_params = frozenset( [ "num_inference_steps", @@ -59,7 +58,7 @@ class QwenControlNetPipelineFastTests(PipelineTesterMixin, unittest.TestCase): ) supports_dduf = False - test_xformers_attention = True + test_xformers_attention = False test_layerwise_casting = True test_group_offloading = True From 7a587349943325e667866971a36996a56fcff143 Mon Sep 17 00:00:00 2001 From: Yao Matrix Date: Tue, 23 Sep 2025 21:01:45 -0700 Subject: [PATCH 022/123] xpu enabling for 4 cases (#12345) Signed-off-by: Yao, Matrix --- .../modular_pipelines/components_manager.py | 13 ++++++++++--- tests/lora/test_lora_layers_hunyuanvideo.py | 1 + 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/src/diffusers/modular_pipelines/components_manager.py b/src/diffusers/modular_pipelines/components_manager.py index f48a227e2e..ed847fa414 100644 --- a/src/diffusers/modular_pipelines/components_manager.py +++ b/src/diffusers/modular_pipelines/components_manager.py @@ -25,6 +25,7 @@ from ..utils import ( is_accelerate_available, logging, ) +from ..utils.torch_utils import get_device if is_accelerate_available(): @@ -161,7 +162,9 @@ class AutoOffloadStrategy: current_module_size = model.get_memory_footprint() - mem_on_device = torch.cuda.mem_get_info(execution_device.index)[0] + device_type = execution_device.type + device_module = getattr(torch, device_type, torch.cuda) + mem_on_device = device_module.mem_get_info(execution_device.index)[0] mem_on_device = mem_on_device - self.memory_reserve_margin if current_module_size < mem_on_device: return [] @@ -301,7 +304,7 @@ class ComponentsManager: cm.add("vae", vae_model, collection="sdxl") # Enable auto offloading - cm.enable_auto_cpu_offload(device="cuda") + cm.enable_auto_cpu_offload() # Retrieve components unet = cm.get_one(name="unet", collection="sdxl") @@ -490,6 +493,8 @@ class ComponentsManager: gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() + if torch.xpu.is_available(): + torch.xpu.empty_cache() # YiYi TODO: rename to search_components for now, may remove this method def search_components( @@ -678,7 +683,7 @@ class ComponentsManager: return get_return_dict(matches, return_dict_with_names) - def enable_auto_cpu_offload(self, device: Union[str, int, torch.device] = "cuda", memory_reserve_margin="3GB"): + def enable_auto_cpu_offload(self, device: Union[str, int, torch.device] = None, memory_reserve_margin="3GB"): """ Enable automatic CPU offloading for all components. @@ -704,6 +709,8 @@ class ComponentsManager: self.disable_auto_cpu_offload() offload_strategy = AutoOffloadStrategy(memory_reserve_margin=memory_reserve_margin) + if device is None: + device = get_device() device = torch.device(device) if device.index is None: device = torch.device(f"{device.type}:{0}") diff --git a/tests/lora/test_lora_layers_hunyuanvideo.py b/tests/lora/test_lora_layers_hunyuanvideo.py index 7ea0f1fcc9..cfd5d3146a 100644 --- a/tests/lora/test_lora_layers_hunyuanvideo.py +++ b/tests/lora/test_lora_layers_hunyuanvideo.py @@ -253,6 +253,7 @@ class HunyuanVideoLoRAIntegrationTests(unittest.TestCase): expected_slices = Expectations( { ("cuda", 7): np.array([0.1013, 0.1924, 0.0078, 0.1021, 0.1929, 0.0078, 0.1023, 0.1919, 0.7402, 0.104, 0.4482, 0.7354, 0.0925, 0.4382, 0.7275, 0.0815]), + ("xpu", 3): np.array([0.1013, 0.1924, 0.0078, 0.1021, 0.1929, 0.0078, 0.1023, 0.1919, 0.7402, 0.104, 0.4482, 0.7354, 0.0925, 0.4382, 0.7275, 0.0815]), } ) # fmt: on From 08c29020dd558899a53137441f908ab668c427d6 Mon Sep 17 00:00:00 2001 From: Yao Matrix Date: Tue, 23 Sep 2025 21:02:06 -0700 Subject: [PATCH 023/123] fix marigold ut case fail on xpu (#12350) Signed-off-by: Yao, Matrix --- .../pipelines/marigold/test_marigold_depth.py | 30 +++++++++++++------ 1 file changed, 21 insertions(+), 9 deletions(-) diff --git a/tests/pipelines/marigold/test_marigold_depth.py b/tests/pipelines/marigold/test_marigold_depth.py index 3e8ccbf5c0..3c85305992 100644 --- a/tests/pipelines/marigold/test_marigold_depth.py +++ b/tests/pipelines/marigold/test_marigold_depth.py @@ -33,6 +33,7 @@ from diffusers import ( ) from ...testing_utils import ( + Expectations, backend_empty_cache, enable_full_determinism, floats_tensor, @@ -356,7 +357,7 @@ class MarigoldDepthPipelineIntegrationTests(unittest.TestCase): match_input_resolution=True, ) - def test_marigold_depth_einstein_f32_cuda_G0_S1_P768_E1_B1_M1(self): + def test_marigold_depth_einstein_f32_accelerator_G0_S1_P768_E1_B1_M1(self): self._test_marigold_depth( is_fp16=False, device=torch_device, @@ -369,7 +370,7 @@ class MarigoldDepthPipelineIntegrationTests(unittest.TestCase): match_input_resolution=True, ) - def test_marigold_depth_einstein_f16_cuda_G0_S1_P768_E1_B1_M1(self): + def test_marigold_depth_einstein_f16_accelerator_G0_S1_P768_E1_B1_M1(self): self._test_marigold_depth( is_fp16=True, device=torch_device, @@ -382,7 +383,7 @@ class MarigoldDepthPipelineIntegrationTests(unittest.TestCase): match_input_resolution=True, ) - def test_marigold_depth_einstein_f16_cuda_G2024_S1_P768_E1_B1_M1(self): + def test_marigold_depth_einstein_f16_accelerator_G2024_S1_P768_E1_B1_M1(self): self._test_marigold_depth( is_fp16=True, device=torch_device, @@ -395,12 +396,23 @@ class MarigoldDepthPipelineIntegrationTests(unittest.TestCase): match_input_resolution=True, ) - def test_marigold_depth_einstein_f16_cuda_G0_S2_P768_E1_B1_M1(self): + def test_marigold_depth_einstein_f16_accelerator_G0_S2_P768_E1_B1_M1(self): + # fmt: off + expected_slices = Expectations( + { + ("cuda", 7): np.array([0.1085, 0.1098, 0.1110, 0.1081, 0.1085, 0.1082, 0.1085, 0.1057, 0.0996]), + ("xpu", 3): np.array([0.1084, 0.1096, 0.1108, 0.1080, 0.1083, 0.1080, + 0.1085, 0.1057, 0.0996]), + } + ) + expected_slice = expected_slices.get_expectation() + # fmt: on + self._test_marigold_depth( is_fp16=True, device=torch_device, generator_seed=0, - expected_slice=np.array([0.1085, 0.1098, 0.1110, 0.1081, 0.1085, 0.1082, 0.1085, 0.1057, 0.0996]), + expected_slice=expected_slice, num_inference_steps=2, processing_resolution=768, ensemble_size=1, @@ -408,7 +420,7 @@ class MarigoldDepthPipelineIntegrationTests(unittest.TestCase): match_input_resolution=True, ) - def test_marigold_depth_einstein_f16_cuda_G0_S1_P512_E1_B1_M1(self): + def test_marigold_depth_einstein_f16_accelerator_G0_S1_P512_E1_B1_M1(self): self._test_marigold_depth( is_fp16=True, device=torch_device, @@ -421,7 +433,7 @@ class MarigoldDepthPipelineIntegrationTests(unittest.TestCase): match_input_resolution=True, ) - def test_marigold_depth_einstein_f16_cuda_G0_S1_P768_E3_B1_M1(self): + def test_marigold_depth_einstein_f16_accelerator_G0_S1_P768_E3_B1_M1(self): self._test_marigold_depth( is_fp16=True, device=torch_device, @@ -435,7 +447,7 @@ class MarigoldDepthPipelineIntegrationTests(unittest.TestCase): match_input_resolution=True, ) - def test_marigold_depth_einstein_f16_cuda_G0_S1_P768_E4_B2_M1(self): + def test_marigold_depth_einstein_f16_accelerator_G0_S1_P768_E4_B2_M1(self): self._test_marigold_depth( is_fp16=True, device=torch_device, @@ -449,7 +461,7 @@ class MarigoldDepthPipelineIntegrationTests(unittest.TestCase): match_input_resolution=True, ) - def test_marigold_depth_einstein_f16_cuda_G0_S1_P512_E1_B1_M0(self): + def test_marigold_depth_einstein_f16_accelerator_G0_S1_P512_E1_B1_M0(self): self._test_marigold_depth( is_fp16=True, device=torch_device, From 043ab2520f6a19fce78e6e060a68dbc947edb9f9 Mon Sep 17 00:00:00 2001 From: Alberto Chimenti <43784986+albchim@users.noreply.github.com> Date: Wed, 24 Sep 2025 11:45:04 +0200 Subject: [PATCH 024/123] Fix WanVACEPipeline to allow prompt to be None and skip encoding step (#12251) Fixed WanVACEPipeline to allow prompt to be None and skip encoding step --- src/diffusers/pipelines/wan/pipeline_wan_vace.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/pipelines/wan/pipeline_wan_vace.py b/src/diffusers/pipelines/wan/pipeline_wan_vace.py index eab1aacfc5..2b1890afec 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan_vace.py +++ b/src/diffusers/pipelines/wan/pipeline_wan_vace.py @@ -795,7 +795,7 @@ class WanVACEPipeline(DiffusionPipeline, WanLoraLoaderMixin): callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs # Simplification of implementation for now - if not isinstance(prompt, str): + if prompt is not None and not isinstance(prompt, str): raise ValueError("Passing a list of prompts is not yet supported. This may be supported in the future.") if num_videos_per_prompt != 1: raise ValueError( From dcb6dd9b7a6c7ddd6875506f40597c0976fd02c5 Mon Sep 17 00:00:00 2001 From: Aryan Date: Wed, 24 Sep 2025 19:03:25 +0530 Subject: [PATCH 025/123] Context Parallel w/ Ring & Ulysses & Unified Attention (#11941) * update * update * add coauthor Co-Authored-By: Dhruv Nair * improve test * handle ip adapter params correctly * fix chroma qkv fusion test * fix fastercache implementation * fix more tests * fight more tests * add back set_attention_backend * update * update * make style * make fix-copies * make ip adapter processor compatible with attention dispatcher * refactor chroma as well * remove rmsnorm assert * minify and deprecate npu/xla processors * update * refactor * refactor; support flash attention 2 with cp * fix * support sage attention with cp * make torch compile compatible * update * refactor * update * refactor * refactor * add ulysses backward * try to make dreambooth script work; accelerator backward not playing well * Revert "try to make dreambooth script work; accelerator backward not playing well" This reverts commit 768d0ea6fa6a305d12df1feda2afae3ec80aa449. * workaround compilation problems with triton when doing all-to-all * support wan * handle backward correctly * support qwen * support ltx * make fix-copies * Update src/diffusers/models/modeling_utils.py Co-authored-by: Dhruv Nair * apply review suggestions * update docs * add explanation * make fix-copies * add docstrings * support passing parallel_config to from_pretrained * apply review suggestions * make style * update * Update docs/source/en/api/parallel.md Co-authored-by: Aryan * up --------- Co-authored-by: Dhruv Nair Co-authored-by: sayakpaul --- docs/source/en/_toctree.yml | 2 + docs/source/en/api/parallel.md | 24 + src/diffusers/__init__.py | 4 + src/diffusers/hooks/__init__.py | 1 + src/diffusers/hooks/context_parallel.py | 297 ++++++ src/diffusers/models/__init__.py | 2 + src/diffusers/models/_modeling_parallel.py | 241 +++++ src/diffusers/models/attention_dispatch.py | 987 +++++++++++++++--- src/diffusers/models/modeling_utils.py | 78 +- .../models/transformers/transformer_bria.py | 8 +- .../models/transformers/transformer_flux.py | 21 +- .../models/transformers/transformer_ltx.py | 15 + .../transformers/transformer_qwenimage.py | 15 + .../transformers/transformer_skyreels_v2.py | 3 + .../models/transformers/transformer_wan.py | 17 + src/diffusers/utils/dummy_pt_objects.py | 30 + 16 files changed, 1571 insertions(+), 174 deletions(-) create mode 100644 docs/source/en/api/parallel.md create mode 100644 src/diffusers/hooks/context_parallel.py create mode 100644 src/diffusers/models/_modeling_parallel.py diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 4879a7bf04..07408cc30e 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -70,6 +70,8 @@ title: Reduce memory usage - local: optimization/speed-memory-optims title: Compiling and offloading quantized models + - local: api/parallel + title: Parallel inference - title: Community optimizations sections: - local: optimization/pruna diff --git a/docs/source/en/api/parallel.md b/docs/source/en/api/parallel.md new file mode 100644 index 0000000000..e38ffe571e --- /dev/null +++ b/docs/source/en/api/parallel.md @@ -0,0 +1,24 @@ + + +# Parallelism + +Parallelism strategies help speed up diffusion transformers by distributing computations across multiple devices, allowing for faster inference/training times. + +## ParallelConfig + +[[autodoc]] ParallelConfig + +## ContextParallelConfig + +[[autodoc]] ContextParallelConfig + +[[autodoc]] hooks.apply_context_parallel diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 741fcd14f2..8867250ded 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -202,6 +202,7 @@ else: "CogView4Transformer2DModel", "ConsisIDTransformer3DModel", "ConsistencyDecoderVAE", + "ContextParallelConfig", "ControlNetModel", "ControlNetUnionModel", "ControlNetXSAdapter", @@ -229,6 +230,7 @@ else: "MultiAdapter", "MultiControlNetModel", "OmniGenTransformer2DModel", + "ParallelConfig", "PixArtTransformer2DModel", "PriorTransformer", "QwenImageControlNetModel", @@ -888,6 +890,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: CogView4Transformer2DModel, ConsisIDTransformer3DModel, ConsistencyDecoderVAE, + ContextParallelConfig, ControlNetModel, ControlNetUnionModel, ControlNetXSAdapter, @@ -915,6 +918,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: MultiAdapter, MultiControlNetModel, OmniGenTransformer2DModel, + ParallelConfig, PixArtTransformer2DModel, PriorTransformer, QwenImageControlNetModel, diff --git a/src/diffusers/hooks/__init__.py b/src/diffusers/hooks/__init__.py index 525a0747da..524a92ea99 100644 --- a/src/diffusers/hooks/__init__.py +++ b/src/diffusers/hooks/__init__.py @@ -16,6 +16,7 @@ from ..utils import is_torch_available if is_torch_available(): + from .context_parallel import apply_context_parallel from .faster_cache import FasterCacheConfig, apply_faster_cache from .first_block_cache import FirstBlockCacheConfig, apply_first_block_cache from .group_offloading import apply_group_offloading diff --git a/src/diffusers/hooks/context_parallel.py b/src/diffusers/hooks/context_parallel.py new file mode 100644 index 0000000000..83406d4969 --- /dev/null +++ b/src/diffusers/hooks/context_parallel.py @@ -0,0 +1,297 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +from dataclasses import dataclass +from typing import Dict, List, Type, Union + +import torch +import torch.distributed._functional_collectives as funcol + +from ..models._modeling_parallel import ( + ContextParallelConfig, + ContextParallelInput, + ContextParallelModelPlan, + ContextParallelOutput, +) +from ..utils import get_logger +from ..utils.torch_utils import unwrap_module +from .hooks import HookRegistry, ModelHook + + +logger = get_logger(__name__) # pylint: disable=invalid-name + +_CONTEXT_PARALLEL_INPUT_HOOK_TEMPLATE = "cp_input---{}" +_CONTEXT_PARALLEL_OUTPUT_HOOK_TEMPLATE = "cp_output---{}" + + +# TODO(aryan): consolidate with ._helpers.TransformerBlockMetadata +@dataclass +class ModuleForwardMetadata: + cached_parameter_indices: Dict[str, int] = None + _cls: Type = None + + def _get_parameter_from_args_kwargs(self, identifier: str, args=(), kwargs=None): + kwargs = kwargs or {} + + if identifier in kwargs: + return kwargs[identifier], True, None + + if self.cached_parameter_indices is not None: + index = self.cached_parameter_indices.get(identifier, None) + if index is None: + raise ValueError(f"Parameter '{identifier}' not found in cached indices.") + return args[index], False, index + + if self._cls is None: + raise ValueError("Model class is not set for metadata.") + + parameters = list(inspect.signature(self._cls.forward).parameters.keys()) + parameters = parameters[1:] # skip `self` + self.cached_parameter_indices = {param: i for i, param in enumerate(parameters)} + + if identifier not in self.cached_parameter_indices: + raise ValueError(f"Parameter '{identifier}' not found in function signature but was requested.") + + index = self.cached_parameter_indices[identifier] + + if index >= len(args): + raise ValueError(f"Expected {index} arguments but got {len(args)}.") + + return args[index], False, index + + +def apply_context_parallel( + module: torch.nn.Module, + parallel_config: ContextParallelConfig, + plan: Dict[str, ContextParallelModelPlan], +) -> None: + """Apply context parallel on a model.""" + logger.debug(f"Applying context parallel with CP mesh: {parallel_config._mesh} and plan: {plan}") + + for module_id, cp_model_plan in plan.items(): + submodule = _get_submodule_by_name(module, module_id) + if not isinstance(submodule, list): + submodule = [submodule] + + logger.debug(f"Applying ContextParallelHook to {module_id=} identifying a total of {len(submodule)} modules") + + for m in submodule: + if isinstance(cp_model_plan, dict): + hook = ContextParallelSplitHook(cp_model_plan, parallel_config) + hook_name = _CONTEXT_PARALLEL_INPUT_HOOK_TEMPLATE.format(module_id) + elif isinstance(cp_model_plan, (ContextParallelOutput, list, tuple)): + if isinstance(cp_model_plan, ContextParallelOutput): + cp_model_plan = [cp_model_plan] + if not all(isinstance(x, ContextParallelOutput) for x in cp_model_plan): + raise ValueError(f"Expected all elements of cp_model_plan to be CPOutput, but got {cp_model_plan}") + hook = ContextParallelGatherHook(cp_model_plan, parallel_config) + hook_name = _CONTEXT_PARALLEL_OUTPUT_HOOK_TEMPLATE.format(module_id) + else: + raise ValueError(f"Unsupported context parallel model plan type: {type(cp_model_plan)}") + registry = HookRegistry.check_if_exists_or_initialize(m) + registry.register_hook(hook, hook_name) + + +def remove_context_parallel(module: torch.nn.Module, plan: Dict[str, ContextParallelModelPlan]) -> None: + for module_id, cp_model_plan in plan.items(): + submodule = _get_submodule_by_name(module, module_id) + if not isinstance(submodule, list): + submodule = [submodule] + + for m in submodule: + registry = HookRegistry.check_if_exists_or_initialize(m) + if isinstance(cp_model_plan, dict): + hook_name = _CONTEXT_PARALLEL_INPUT_HOOK_TEMPLATE.format(module_id) + elif isinstance(cp_model_plan, (ContextParallelOutput, list, tuple)): + hook_name = _CONTEXT_PARALLEL_OUTPUT_HOOK_TEMPLATE.format(module_id) + else: + raise ValueError(f"Unsupported context parallel model plan type: {type(cp_model_plan)}") + registry.remove_hook(hook_name) + + +class ContextParallelSplitHook(ModelHook): + def __init__(self, metadata: ContextParallelModelPlan, parallel_config: ContextParallelConfig) -> None: + super().__init__() + self.metadata = metadata + self.parallel_config = parallel_config + self.module_forward_metadata = None + + def initialize_hook(self, module): + cls = unwrap_module(module).__class__ + self.module_forward_metadata = ModuleForwardMetadata(_cls=cls) + return module + + def pre_forward(self, module, *args, **kwargs): + args_list = list(args) + + for name, cpm in self.metadata.items(): + if isinstance(cpm, ContextParallelInput) and cpm.split_output: + continue + + # Maybe the parameter was passed as a keyword argument + input_val, is_kwarg, index = self.module_forward_metadata._get_parameter_from_args_kwargs( + name, args_list, kwargs + ) + + if input_val is None: + continue + + # The input_val may be a tensor or list/tuple of tensors. In certain cases, user may specify to shard + # the output instead of input for a particular layer by setting split_output=True + if isinstance(input_val, torch.Tensor): + input_val = self._prepare_cp_input(input_val, cpm) + elif isinstance(input_val, (list, tuple)): + if len(input_val) != len(cpm): + raise ValueError( + f"Expected input model plan to have {len(input_val)} elements, but got {len(cpm)}." + ) + sharded_input_val = [] + for i, x in enumerate(input_val): + if torch.is_tensor(x) and not cpm[i].split_output: + x = self._prepare_cp_input(x, cpm[i]) + sharded_input_val.append(x) + input_val = sharded_input_val + else: + raise ValueError(f"Unsupported input type: {type(input_val)}") + + if is_kwarg: + kwargs[name] = input_val + elif index is not None and index < len(args_list): + args_list[index] = input_val + else: + raise ValueError( + f"An unexpected error occurred while processing the input '{name}'. Please open an " + f"issue at https://github.com/huggingface/diffusers/issues and provide a minimal reproducible " + f"example along with the full stack trace." + ) + + return tuple(args_list), kwargs + + def post_forward(self, module, output): + is_tensor = isinstance(output, torch.Tensor) + is_tensor_list = isinstance(output, (list, tuple)) and all(isinstance(x, torch.Tensor) for x in output) + + if not is_tensor and not is_tensor_list: + raise ValueError(f"Expected output to be a tensor or a list/tuple of tensors, but got {type(output)}.") + + output = [output] if is_tensor else list(output) + for index, cpm in self.metadata.items(): + if not isinstance(cpm, ContextParallelInput) or not cpm.split_output: + continue + if index >= len(output): + raise ValueError(f"Index {index} out of bounds for output of length {len(output)}.") + current_output = output[index] + current_output = self._prepare_cp_input(current_output, cpm) + output[index] = current_output + + return output[0] if is_tensor else tuple(output) + + def _prepare_cp_input(self, x: torch.Tensor, cp_input: ContextParallelInput) -> torch.Tensor: + if cp_input.expected_dims is not None and x.dim() != cp_input.expected_dims: + raise ValueError( + f"Expected input tensor to have {cp_input.expected_dims} dimensions, but got {x.dim()} dimensions." + ) + return EquipartitionSharder.shard(x, cp_input.split_dim, self.parallel_config._flattened_mesh) + + +class ContextParallelGatherHook(ModelHook): + def __init__(self, metadata: ContextParallelModelPlan, parallel_config: ContextParallelConfig) -> None: + super().__init__() + self.metadata = metadata + self.parallel_config = parallel_config + + def post_forward(self, module, output): + is_tensor = isinstance(output, torch.Tensor) + + if is_tensor: + output = [output] + elif not (isinstance(output, (list, tuple)) and all(isinstance(x, torch.Tensor) for x in output)): + raise ValueError(f"Expected output to be a tensor or a list/tuple of tensors, but got {type(output)}.") + + output = list(output) + + if len(output) != len(self.metadata): + raise ValueError(f"Expected output to have {len(self.metadata)} elements, but got {len(output)}.") + + for i, cpm in enumerate(self.metadata): + if cpm is None: + continue + output[i] = EquipartitionSharder.unshard(output[i], cpm.gather_dim, self.parallel_config._flattened_mesh) + + return output[0] if is_tensor else tuple(output) + + +class AllGatherFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, tensor, dim, group): + ctx.dim = dim + ctx.group = group + ctx.world_size = torch.distributed.get_world_size(group) + ctx.rank = torch.distributed.get_rank(group) + return funcol.all_gather_tensor(tensor, dim, group=group) + + @staticmethod + def backward(ctx, grad_output): + grad_chunks = torch.chunk(grad_output, ctx.world_size, dim=ctx.dim) + return grad_chunks[ctx.rank], None, None + + +class EquipartitionSharder: + @classmethod + def shard(cls, tensor: torch.Tensor, dim: int, mesh: torch.distributed.device_mesh.DeviceMesh) -> torch.Tensor: + # NOTE: the following assertion does not have to be true in general. We simply enforce it for now + # because the alternate case has not yet been tested/required for any model. + assert tensor.size()[dim] % mesh.size() == 0, ( + "Tensor size along dimension to be sharded must be divisible by mesh size" + ) + + # The following is not fullgraph compatible with Dynamo (fails in DeviceMesh.get_rank) + # return tensor.chunk(mesh.size(), dim=dim)[mesh.get_rank()] + + return tensor.chunk(mesh.size(), dim=dim)[torch.distributed.get_rank(mesh.get_group())] + + @classmethod + def unshard(cls, tensor: torch.Tensor, dim: int, mesh: torch.distributed.device_mesh.DeviceMesh) -> torch.Tensor: + tensor = tensor.contiguous() + tensor = AllGatherFunction.apply(tensor, dim, mesh.get_group()) + return tensor + + +def _get_submodule_by_name(model: torch.nn.Module, name: str) -> Union[torch.nn.Module, List[torch.nn.Module]]: + if name.count("*") > 1: + raise ValueError("Wildcard '*' can only be used once in the name") + return _find_submodule_by_name(model, name) + + +def _find_submodule_by_name(model: torch.nn.Module, name: str) -> Union[torch.nn.Module, List[torch.nn.Module]]: + if name == "": + return model + first_atom, remaining_name = name.split(".", 1) if "." in name else (name, "") + if first_atom == "*": + if not isinstance(model, torch.nn.ModuleList): + raise ValueError("Wildcard '*' can only be used with ModuleList") + submodules = [] + for submodule in model: + subsubmodules = _find_submodule_by_name(submodule, remaining_name) + if not isinstance(subsubmodules, list): + subsubmodules = [subsubmodules] + submodules.extend(subsubmodules) + return submodules + else: + if hasattr(model, first_atom): + submodule = getattr(model, first_atom) + return _find_submodule_by_name(submodule, remaining_name) + else: + raise ValueError(f"'{first_atom}' is not a submodule of '{model.__class__.__name__}'") diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index 49ac2a1c56..457f70448a 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -25,6 +25,7 @@ from ..utils import ( _import_structure = {} if is_torch_available(): + _import_structure["_modeling_parallel"] = ["ContextParallelConfig", "ParallelConfig"] _import_structure["adapter"] = ["MultiAdapter", "T2IAdapter"] _import_structure["attention_dispatch"] = ["AttentionBackendName", "attention_backend"] _import_structure["auto_model"] = ["AutoModel"] @@ -119,6 +120,7 @@ if is_flax_available(): if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: if is_torch_available(): + from ._modeling_parallel import ContextParallelConfig, ParallelConfig from .adapter import MultiAdapter, T2IAdapter from .attention_dispatch import AttentionBackendName, attention_backend from .auto_model import AutoModel diff --git a/src/diffusers/models/_modeling_parallel.py b/src/diffusers/models/_modeling_parallel.py new file mode 100644 index 0000000000..2a1d2cc6ce --- /dev/null +++ b/src/diffusers/models/_modeling_parallel.py @@ -0,0 +1,241 @@ +# 🚨🚨🚨 Experimental parallelism support for Diffusers 🚨🚨🚨 +# Experimental changes are subject to change and APIs may break without warning. + +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, Union + +import torch + +from ..utils import get_logger + + +if TYPE_CHECKING: + pass + + +logger = get_logger(__name__) # pylint: disable=invalid-name + + +# TODO(aryan): add support for the following: +# - Unified Attention +# - More dispatcher attention backends +# - CFG/Data Parallel +# - Tensor Parallel + + +@dataclass +class ContextParallelConfig: + """ + Configuration for context parallelism. + + Args: + ring_degree (`int`, *optional*, defaults to `1`): + Number of devices to use for ring attention within a context parallel region. Must be a divisor of the + total number of devices in the context parallel mesh. + ulysses_degree (`int`, *optional*, defaults to `1`): + Number of devices to use for ulysses attention within a context parallel region. Must be a divisor of the + total number of devices in the context parallel mesh. + convert_to_fp32 (`bool`, *optional*, defaults to `True`): + Whether to convert output and LSE to float32 for ring attention numerical stability. + rotate_method (`str`, *optional*, defaults to `"allgather"`): + Method to use for rotating key/value states across devices in ring attention. Currently, only `"allgather"` + is supported. + + """ + + ring_degree: Optional[int] = None + ulysses_degree: Optional[int] = None + convert_to_fp32: bool = True + # TODO: support alltoall + rotate_method: Literal["allgather", "alltoall"] = "allgather" + + _rank: int = None + _world_size: int = None + _device: torch.device = None + _mesh: torch.distributed.device_mesh.DeviceMesh = None + _flattened_mesh: torch.distributed.device_mesh.DeviceMesh = None + _ring_mesh: torch.distributed.device_mesh.DeviceMesh = None + _ulysses_mesh: torch.distributed.device_mesh.DeviceMesh = None + _ring_local_rank: int = None + _ulysses_local_rank: int = None + + def __post_init__(self): + if self.ring_degree is None: + self.ring_degree = 1 + if self.ulysses_degree is None: + self.ulysses_degree = 1 + + def setup(self, rank: int, world_size: int, device: torch.device, mesh: torch.distributed.device_mesh.DeviceMesh): + self._rank = rank + self._world_size = world_size + self._device = device + self._mesh = mesh + if self.ring_degree is None: + self.ring_degree = 1 + if self.ulysses_degree is None: + self.ulysses_degree = 1 + if self.rotate_method != "allgather": + raise NotImplementedError( + f"Only rotate_method='allgather' is supported for now, but got {self.rotate_method}." + ) + if self._flattened_mesh is None: + self._flattened_mesh = self._mesh._flatten() + if self._ring_mesh is None: + self._ring_mesh = self._mesh["ring"] + if self._ulysses_mesh is None: + self._ulysses_mesh = self._mesh["ulysses"] + if self._ring_local_rank is None: + self._ring_local_rank = self._ring_mesh.get_local_rank() + if self._ulysses_local_rank is None: + self._ulysses_local_rank = self._ulysses_mesh.get_local_rank() + + +@dataclass +class ParallelConfig: + """ + Configuration for applying different parallelisms. + + Args: + context_parallel_config (`ContextParallelConfig`, *optional*): + Configuration for context parallelism. + """ + + context_parallel_config: Optional[ContextParallelConfig] = None + + _rank: int = None + _world_size: int = None + _device: torch.device = None + _cp_mesh: torch.distributed.device_mesh.DeviceMesh = None + + def setup( + self, + rank: int, + world_size: int, + device: torch.device, + *, + cp_mesh: Optional[torch.distributed.device_mesh.DeviceMesh] = None, + ): + self._rank = rank + self._world_size = world_size + self._device = device + self._cp_mesh = cp_mesh + if self.context_parallel_config is not None: + self.context_parallel_config.setup(rank, world_size, device, cp_mesh) + + +@dataclass(frozen=True) +class ContextParallelInput: + """ + Configuration for splitting an input tensor across context parallel region. + + Args: + split_dim (`int`): + The dimension along which to split the tensor. + expected_dims (`int`, *optional*): + The expected number of dimensions of the tensor. If provided, a check will be performed to ensure that the + tensor has the expected number of dimensions before splitting. + split_output (`bool`, *optional*, defaults to `False`): + Whether to split the output tensor of the layer along the given `split_dim` instead of the input tensor. + This is useful for layers whose outputs should be split after it does some preprocessing on the inputs (ex: + RoPE). + """ + + split_dim: int + expected_dims: Optional[int] = None + split_output: bool = False + + def __repr__(self): + return f"ContextParallelInput(split_dim={self.split_dim}, expected_dims={self.expected_dims}, split_output={self.split_output})" + + +@dataclass(frozen=True) +class ContextParallelOutput: + """ + Configuration for gathering an output tensor across context parallel region. + + Args: + gather_dim (`int`): + The dimension along which to gather the tensor. + expected_dims (`int`, *optional*): + The expected number of dimensions of the tensor. If provided, a check will be performed to ensure that the + tensor has the expected number of dimensions before gathering. + """ + + gather_dim: int + expected_dims: Optional[int] = None + + def __repr__(self): + return f"ContextParallelOutput(gather_dim={self.gather_dim}, expected_dims={self.expected_dims})" + + +# A dictionary where keys denote the input to be split across context parallel region, and the +# value denotes the sharding configuration. +# If the key is a string, it denotes the name of the parameter in the forward function. +# If the key is an integer, split_output must be set to True, and it denotes the index of the output +# to be split across context parallel region. +ContextParallelInputType = Dict[ + Union[str, int], Union[ContextParallelInput, List[ContextParallelInput], Tuple[ContextParallelInput, ...]] +] + +# A dictionary where keys denote the output to be gathered across context parallel region, and the +# value denotes the gathering configuration. +ContextParallelOutputType = Union[ + ContextParallelOutput, List[ContextParallelOutput], Tuple[ContextParallelOutput, ...] +] + +# A dictionary where keys denote the module id, and the value denotes how the inputs/outputs of +# the module should be split/gathered across context parallel region. +ContextParallelModelPlan = Dict[str, Union[ContextParallelInputType, ContextParallelOutputType]] + + +# Example of a ContextParallelModelPlan (QwenImageTransformer2DModel): +# +# Each model should define a _cp_plan attribute that contains information on how to shard/gather +# tensors at different stages of the forward: +# +# ```python +# _cp_plan = { +# "": { +# "hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False), +# "encoder_hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False), +# "encoder_hidden_states_mask": ContextParallelInput(split_dim=1, expected_dims=2, split_output=False), +# }, +# "pos_embed": { +# 0: ContextParallelInput(split_dim=0, expected_dims=2, split_output=True), +# 1: ContextParallelInput(split_dim=0, expected_dims=2, split_output=True), +# }, +# "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3), +# } +# ``` +# +# The dictionary is a set of module names mapped to their respective CP plan. The inputs/outputs of layers will be +# split/gathered according to this at the respective module level. Here, the following happens: +# - "": +# we specify that we want to split the various inputs across the sequence dim in the pre-forward hook (i.e. before +# the actual forward logic of the QwenImageTransformer2DModel is run, we will splitthe inputs) +# - "pos_embed": +# we specify that we want to split the outputs of the RoPE layer. Since there are two outputs (imag & text freqs), +# we can individually specify how they should be split +# - "proj_out": +# before returning to the user, we gather the entire sequence on each rank in the post-forward hook (after the linear +# layer forward has run). +# +# ContextParallelInput: +# specifies how to split the input tensor in the pre-forward or post-forward hook of the layer it is attached to +# +# ContextParallelOutput: +# specifies how to gather the input tensor in the post-forward hook in the layer it is attached to diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index f71be7c8ec..0a2ad68123 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -17,9 +17,10 @@ import functools import inspect import math from enum import Enum -from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Tuple, Union import torch +import torch.distributed._functional_collectives as funcol from ..utils import ( get_logger, @@ -39,6 +40,9 @@ from ..utils import ( from ..utils.constants import DIFFUSERS_ATTN_BACKEND, DIFFUSERS_ATTN_CHECKS, DIFFUSERS_ENABLE_HUB_KERNELS +if TYPE_CHECKING: + from ._modeling_parallel import ParallelConfig + _REQUIRED_FLASH_VERSION = "2.6.3" _REQUIRED_SAGE_VERSION = "2.1.1" _REQUIRED_FLEX_VERSION = "2.5.0" @@ -56,9 +60,12 @@ _CAN_USE_XFORMERS_ATTN = is_xformers_available() and is_xformers_version(">=", _ if _CAN_USE_FLASH_ATTN: from flash_attn import flash_attn_func, flash_attn_varlen_func + from flash_attn.flash_attn_interface import _wrapped_flash_attn_backward, _wrapped_flash_attn_forward else: flash_attn_func = None flash_attn_varlen_func = None + _wrapped_flash_attn_backward = None + _wrapped_flash_attn_forward = None if _CAN_USE_FLASH_ATTN_3: @@ -197,17 +204,24 @@ class _AttentionBackendRegistry: _backends = {} _constraints = {} _supported_arg_names = {} + _supports_context_parallel = {} _active_backend = AttentionBackendName(DIFFUSERS_ATTN_BACKEND) _checks_enabled = DIFFUSERS_ATTN_CHECKS @classmethod - def register(cls, backend: AttentionBackendName, constraints: Optional[List[Callable]] = None): + def register( + cls, + backend: AttentionBackendName, + constraints: Optional[List[Callable]] = None, + supports_context_parallel: bool = False, + ): logger.debug(f"Registering attention backend: {backend} with constraints: {constraints}") def decorator(func): cls._backends[backend] = func cls._constraints[backend] = constraints or [] cls._supported_arg_names[backend] = set(inspect.signature(func).parameters.keys()) + cls._supports_context_parallel[backend] = supports_context_parallel return func return decorator @@ -220,6 +234,17 @@ class _AttentionBackendRegistry: def list_backends(cls): return list(cls._backends.keys()) + @classmethod + def _is_context_parallel_enabled( + cls, backend: AttentionBackendName, parallel_config: Optional["ParallelConfig"] + ) -> bool: + supports_context_parallel = backend in cls._supports_context_parallel + is_degree_greater_than_1 = parallel_config is not None and ( + parallel_config.context_parallel_config.ring_degree > 1 + or parallel_config.context_parallel_config.ulysses_degree > 1 + ) + return supports_context_parallel and is_degree_greater_than_1 + @contextlib.contextmanager def attention_backend(backend: Union[str, AttentionBackendName] = AttentionBackendName.NATIVE): @@ -253,6 +278,7 @@ def dispatch_attention_fn( attention_kwargs: Optional[Dict[str, Any]] = None, *, backend: Optional[AttentionBackendName] = None, + parallel_config: Optional["ParallelConfig"] = None, ) -> torch.Tensor: attention_kwargs = attention_kwargs or {} @@ -264,6 +290,14 @@ def dispatch_attention_fn( backend_name = AttentionBackendName(backend) backend_fn = _AttentionBackendRegistry._backends.get(backend_name) + if parallel_config is not None and not _AttentionBackendRegistry._is_context_parallel_enabled( + backend_name, parallel_config + ): + raise ValueError( + f"Backend {backend_name} either does not support context parallelism or context parallelism " + f"was enabled with a world size of 1." + ) + kwargs = { "query": query, "key": key, @@ -273,6 +307,7 @@ def dispatch_attention_fn( "is_causal": is_causal, "scale": scale, **attention_kwargs, + "_parallel_config": parallel_config, } if is_torch_version(">=", "2.5.0"): kwargs["enable_gqa"] = enable_gqa @@ -521,22 +556,621 @@ def _flex_attention_causal_mask_mod(batch_idx, head_idx, q_idx, kv_idx): # Registrations are required for fullgraph tracing compatibility # TODO: this is only required because the beta release FA3 does not have it. There is a PR adding # this but it was never merged: https://github.com/Dao-AILab/flash-attention/pull/1590 - - -@_custom_op("flash_attn_3::_flash_attn_forward", mutates_args=(), device_types="cuda") -def _wrapped_flash_attn_3_original( - query: torch.Tensor, key: torch.Tensor, value: torch.Tensor +@_custom_op("_diffusers_flash_attn_3::_flash_attn_forward", mutates_args=(), device_types="cuda") +def _wrapped_flash_attn_3( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + softmax_scale: Optional[float] = None, + causal: bool = False, + qv: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + attention_chunk: int = 0, + softcap: float = 0.0, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + deterministic: bool = False, + sm_margin: int = 0, ) -> Tuple[torch.Tensor, torch.Tensor]: - out, lse = flash_attn_3_func(query, key, value) + # Hardcoded for now because pytorch does not support tuple/int type hints + window_size = (-1, -1) + out, lse, *_ = flash_attn_3_func( + q=q, + k=k, + v=v, + softmax_scale=softmax_scale, + causal=causal, + qv=qv, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + deterministic=deterministic, + sm_margin=sm_margin, + ) lse = lse.permute(0, 2, 1) return out, lse -@_register_fake("flash_attn_3::_flash_attn_forward") -def _(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - batch_size, seq_len, num_heads, head_dim = query.shape +@_register_fake("_diffusers_flash_attn_3::_flash_attn_forward") +def _( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + softmax_scale: Optional[float] = None, + causal: bool = False, + qv: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + attention_chunk: int = 0, + softcap: float = 0.0, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + deterministic: bool = False, + sm_margin: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor]: + window_size = (-1, -1) # noqa: F841 + # A lot of the parameters here are not yet used in any way within diffusers. + # We can safely ignore for now and keep the fake op shape propagation simple. + batch_size, seq_len, num_heads, head_dim = q.shape lse_shape = (batch_size, seq_len, num_heads) - return torch.empty_like(query), query.new_empty(lse_shape) + return torch.empty_like(q), q.new_empty(lse_shape) + + +# ===== Helper functions to use attention backends with templated CP autograd functions ===== + + +# https://github.com/pytorch/pytorch/blob/8904ba638726f8c9a5aff5977c4aa76c9d2edfa6/aten/src/ATen/native/native_functions.yaml#L14958 +# forward declaration: +# aten::_scaled_dot_product_cudnn_attention(Tensor query, Tensor key, Tensor value, Tensor? attn_bias, bool compute_log_sumexp, float dropout_p=0., bool is_causal=False, bool return_debug_mask=False, *, float? scale=None) -> (Tensor output, Tensor logsumexp, Tensor cum_seq_q, Tensor cum_seq_k, SymInt max_q, SymInt max_k, Tensor philox_seed, Tensor philox_offset, Tensor debug_attn_mask) +def _cudnn_attention_forward_op( + ctx: torch.autograd.function.FunctionCtx, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: Optional[torch.Tensor] = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: Optional[float] = None, + enable_gqa: bool = False, + return_lse: bool = False, + _save_ctx: bool = True, + _parallel_config: Optional["ParallelConfig"] = None, +): + if enable_gqa: + raise ValueError("`enable_gqa` is not yet supported for cuDNN attention.") + + tensors_to_save = () + + # Contiguous is a must here! Calling cuDNN backend with aten ops produces incorrect results + # if the input tensors are not contiguous. + query = query.transpose(1, 2).contiguous() + key = key.transpose(1, 2).contiguous() + value = value.transpose(1, 2).contiguous() + tensors_to_save += (query, key, value) + + out, lse, cum_seq_q, cum_seq_k, max_q, max_k, philox_seed, philox_offset, debug_attn_mask = ( + torch.ops.aten._scaled_dot_product_cudnn_attention( + query=query, + key=key, + value=value, + attn_bias=attn_mask, + compute_log_sumexp=return_lse, + dropout_p=dropout_p, + is_causal=is_causal, + return_debug_mask=False, + scale=scale, + ) + ) + + tensors_to_save += (out, lse, cum_seq_q, cum_seq_k, philox_seed, philox_offset) + if _save_ctx: + ctx.save_for_backward(*tensors_to_save) + ctx.dropout_p = dropout_p + ctx.is_causal = is_causal + ctx.scale = scale + ctx.attn_mask = attn_mask + ctx.max_q = max_q + ctx.max_k = max_k + + out = out.transpose(1, 2).contiguous() + if lse is not None: + lse = lse.transpose(1, 2).contiguous() + return (out, lse) if return_lse else out + + +# backward declaration: +# aten::_scaled_dot_product_cudnn_attention_backward(Tensor grad_out, Tensor query, Tensor key, Tensor value, Tensor out, Tensor logsumexp, Tensor philox_seed, Tensor philox_offset, Tensor attn_bias, Tensor cum_seq_q, Tensor cum_seq_k, SymInt max_q, SymInt max_k, float dropout_p, bool is_causal, *, float? scale=None) -> (Tensor, Tensor, Tensor) +def _cudnn_attention_backward_op( + ctx: torch.autograd.function.FunctionCtx, + grad_out: torch.Tensor, + *args, + **kwargs, +): + query, key, value, out, lse, cum_seq_q, cum_seq_k, philox_seed, philox_offset = ctx.saved_tensors + + grad_out = grad_out.transpose(1, 2).contiguous() + key = key.transpose(1, 2).contiguous() + value = value.transpose(1, 2).contiguous() + + # Cannot pass first 5 arguments as kwargs because: https://github.com/pytorch/pytorch/blob/d26ca5de058dbcf56ac52bb43e84dd98df2ace97/torch/_dynamo/variables/torch.py#L1341 + grad_query, grad_key, grad_value = torch.ops.aten._scaled_dot_product_cudnn_attention_backward( + grad_out, + query, + key, + value, + out, + logsumexp=lse, + philox_seed=philox_seed, + philox_offset=philox_offset, + attn_bias=ctx.attn_mask, + cum_seq_q=cum_seq_q, + cum_seq_k=cum_seq_k, + max_q=ctx.max_q, + max_k=ctx.max_k, + dropout_p=ctx.dropout_p, + is_causal=ctx.is_causal, + scale=ctx.scale, + ) + grad_query, grad_key, grad_value = (x.transpose(1, 2).contiguous() for x in (grad_query, grad_key, grad_value)) + + return grad_query, grad_key, grad_value + + +# Adapted from: https://github.com/Dao-AILab/flash-attention/blob/fd2fc9d85c8e54e5c20436465bca709bc1a6c5a1/flash_attn/flash_attn_interface.py#L807 +def _flash_attention_forward_op( + ctx: torch.autograd.function.FunctionCtx, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: Optional[torch.Tensor] = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: Optional[float] = None, + enable_gqa: bool = False, + return_lse: bool = False, + _save_ctx: bool = True, + _parallel_config: Optional["ParallelConfig"] = None, +): + if attn_mask is not None: + raise ValueError("`attn_mask` is not yet supported for flash-attn 2.") + if enable_gqa: + raise ValueError("`enable_gqa` is not yet supported for flash-attn 2.") + + # Hardcoded for now + window_size = (-1, -1) + softcap = 0.0 + alibi_slopes = None + deterministic = False + grad_enabled = any(x.requires_grad for x in (query, key, value)) + + if scale is None: + scale = query.shape[-1] ** (-0.5) + + # flash-attn only returns LSE if dropout_p > 0. So, we need to workaround. + if grad_enabled or (_parallel_config is not None and _parallel_config.context_parallel_config._world_size > 1): + dropout_p = dropout_p if dropout_p > 0 else 1e-30 + + with torch.set_grad_enabled(grad_enabled): + out, lse, S_dmask, rng_state = _wrapped_flash_attn_forward( + query, + key, + value, + dropout_p, + scale, + is_causal, + window_size[0], + window_size[1], + softcap, + alibi_slopes, + return_lse, + ) + lse = lse.permute(0, 2, 1) + + if _save_ctx: + ctx.save_for_backward(query, key, value, out, lse, rng_state) + ctx.dropout_p = dropout_p + ctx.scale = scale + ctx.is_causal = is_causal + ctx.window_size = window_size + ctx.softcap = softcap + ctx.alibi_slopes = alibi_slopes + ctx.deterministic = deterministic + + return (out, lse) if return_lse else out + + +def _flash_attention_backward_op( + ctx: torch.autograd.function.FunctionCtx, + grad_out: torch.Tensor, + *args, + **kwargs, +): + query, key, value, out, lse, rng_state = ctx.saved_tensors + grad_query, grad_key, grad_value = torch.empty_like(query), torch.empty_like(key), torch.empty_like(value) + + lse_d = _wrapped_flash_attn_backward( # noqa: F841 + grad_out, + query, + key, + value, + out, + lse, + grad_query, + grad_key, + grad_value, + ctx.dropout_p, + ctx.scale, + ctx.is_causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.alibi_slopes, + ctx.deterministic, + rng_state, + ) + + # Head dimension may have been padded + grad_query = grad_query[..., : grad_out.shape[-1]] + grad_key = grad_key[..., : grad_out.shape[-1]] + grad_value = grad_value[..., : grad_out.shape[-1]] + + return grad_query, grad_key, grad_value + + +def _sage_attention_forward_op( + ctx: torch.autograd.function.FunctionCtx, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: Optional[torch.Tensor] = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: Optional[float] = None, + enable_gqa: bool = False, + return_lse: bool = False, + _save_ctx: bool = True, + _parallel_config: Optional["ParallelConfig"] = None, +): + if attn_mask is not None: + raise ValueError("`attn_mask` is not yet supported for Sage attention.") + if dropout_p > 0.0: + raise ValueError("`dropout_p` is not yet supported for Sage attention.") + if enable_gqa: + raise ValueError("`enable_gqa` is not yet supported for Sage attention.") + + out = sageattn( + q=query, + k=key, + v=value, + tensor_layout="NHD", + is_causal=is_causal, + sm_scale=scale, + return_lse=return_lse, + ) + lse = None + if return_lse: + out, lse, *_ = out + lse = lse.permute(0, 2, 1) + + return (out, lse) if return_lse else out + + +def _sage_attention_backward_op( + ctx: torch.autograd.function.FunctionCtx, + grad_out: torch.Tensor, + *args, +): + raise NotImplementedError("Backward pass is not implemented for Sage attention.") + + +# ===== Context parallel ===== + + +# Reference: +# - https://github.com/pytorch/pytorch/blob/f58a680d09e13658a52c6ba05c63c15759846bcc/torch/distributed/_functional_collectives.py#L827 +# - https://github.com/pytorch/pytorch/blob/f58a680d09e13658a52c6ba05c63c15759846bcc/torch/distributed/_functional_collectives.py#L246 +# For fullgraph=True tracing compatibility (since FakeTensor does not have a `wait` method): +def _wait_tensor(tensor): + if isinstance(tensor, funcol.AsyncCollectiveTensor): + tensor = tensor.wait() + return tensor + + +def _all_to_all_single(x: torch.Tensor, group) -> torch.Tensor: + shape = x.shape + # HACK: We need to flatten because despite making tensors contiguous, torch single-file-ization + # to benchmark triton codegen fails somewhere: + # buf25 = torch.ops._c10d_functional.all_to_all_single.default(buf24, [1, 1], [1, 1], '3') + # ValueError: Tensors must be contiguous + x = x.flatten() + x = funcol.all_to_all_single(x, None, None, group) + x = x.reshape(shape) + x = _wait_tensor(x) + return x + + +class TemplatedRingAttention(torch.autograd.Function): + @staticmethod + def forward( + ctx: torch.autograd.function.FunctionCtx, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: Optional[torch.Tensor], + dropout_p: float, + is_causal: bool, + scale: Optional[float], + enable_gqa: bool, + return_lse: bool, + forward_op, + backward_op, + _parallel_config: Optional["ParallelConfig"] = None, + ): + ring_mesh = _parallel_config.context_parallel_config._ring_mesh + rank = _parallel_config.context_parallel_config._ring_local_rank + world_size = _parallel_config.context_parallel_config.ring_degree + next_rank = (rank + 1) % world_size + prev_out = prev_lse = None + + ctx.forward_op = forward_op + ctx.backward_op = backward_op + ctx.q_shape = query.shape + ctx.kv_shape = key.shape + ctx._parallel_config = _parallel_config + + kv_buffer = torch.cat([key.flatten(), value.flatten()]).contiguous() + kv_buffer = funcol.all_gather_tensor(kv_buffer, gather_dim=0, group=ring_mesh.get_group()) + kv_buffer = kv_buffer.chunk(world_size) + + for i in range(world_size): + if i > 0: + kv = kv_buffer[next_rank] + key_numel = key.numel() + key = kv[:key_numel].reshape_as(key) + value = kv[key_numel:].reshape_as(value) + next_rank = (next_rank + 1) % world_size + + out, lse = forward_op( + ctx, + query, + key, + value, + attn_mask, + dropout_p, + is_causal, + scale, + enable_gqa, + True, + _save_ctx=i == 0, + _parallel_config=_parallel_config, + ) + + if _parallel_config.context_parallel_config.convert_to_fp32: + out = out.to(torch.float32) + lse = lse.to(torch.float32) + + lse = lse.unsqueeze(-1) + if prev_out is not None: + out = prev_out - torch.nn.functional.sigmoid(lse - prev_lse) * (prev_out - out) + lse = prev_lse - torch.nn.functional.logsigmoid(prev_lse - lse) + prev_out = out + prev_lse = lse + + out = out.to(query.dtype) + lse = lse.squeeze(-1) + + return (out, lse) if return_lse else out + + @staticmethod + def backward( + ctx: torch.autograd.function.FunctionCtx, + grad_out: torch.Tensor, + *args, + ): + ring_mesh = ctx._parallel_config.context_parallel_config._ring_mesh + rank = ctx._parallel_config.context_parallel_config._ring_local_rank + world_size = ctx._parallel_config.context_parallel_config.ring_degree + next_rank = (rank + 1) % world_size + next_ranks = list(range(1, world_size)) + [0] + + accum_dtype = torch.float32 if ctx._parallel_config.context_parallel_config.convert_to_fp32 else grad_out.dtype + grad_query = torch.zeros(ctx.q_shape, dtype=accum_dtype, device=grad_out.device) + grad_key = torch.zeros(ctx.kv_shape, dtype=accum_dtype, device=grad_out.device) + grad_value = torch.zeros(ctx.kv_shape, dtype=accum_dtype, device=grad_out.device) + next_grad_kv = None + + query, key, value, *_ = ctx.saved_tensors + kv_buffer = torch.cat([key.flatten(), value.flatten()]).contiguous() + kv_buffer = funcol.all_gather_tensor(kv_buffer, gather_dim=0, group=ring_mesh.get_group()) + kv_buffer = kv_buffer.chunk(world_size) + + for i in range(world_size): + if i > 0: + kv = kv_buffer[next_rank] + key_numel = key.numel() + key = kv[:key_numel].reshape_as(key) + value = kv[key_numel:].reshape_as(value) + next_rank = (next_rank + 1) % world_size + + grad_query_op, grad_key_op, grad_value_op, *_ = ctx.backward_op(ctx, grad_out) + + if i > 0: + grad_kv_buffer = _wait_tensor(next_grad_kv) + grad_key_numel = grad_key.numel() + grad_key = grad_kv_buffer[:grad_key_numel].reshape_as(grad_key) + grad_value = grad_kv_buffer[grad_key_numel:].reshape_as(grad_value) + + grad_query += grad_query_op + grad_key += grad_key_op + grad_value += grad_value_op + + if i < world_size - 1: + grad_kv_buffer = torch.cat([grad_key.flatten(), grad_value.flatten()]).contiguous() + next_grad_kv = funcol.permute_tensor(grad_kv_buffer, next_ranks, group=ring_mesh.get_group()) + + grad_query, grad_key, grad_value = (x.to(grad_out.dtype) for x in (grad_query, grad_key, grad_value)) + + return grad_query, grad_key, grad_value, None, None, None, None, None, None, None, None + + +class TemplatedUlyssesAttention(torch.autograd.Function): + @staticmethod + def forward( + ctx: torch.autograd.function.FunctionCtx, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: Optional[torch.Tensor], + dropout_p: float, + is_causal: bool, + scale: Optional[float], + enable_gqa: bool, + return_lse: bool, + forward_op, + backward_op, + _parallel_config: Optional["ParallelConfig"] = None, + ): + ulysses_mesh = _parallel_config.context_parallel_config._ulysses_mesh + world_size = _parallel_config.context_parallel_config.ulysses_degree + group = ulysses_mesh.get_group() + + ctx.forward_op = forward_op + ctx.backward_op = backward_op + ctx._parallel_config = _parallel_config + + B, S_Q_LOCAL, H, D = query.shape + _, S_KV_LOCAL, _, _ = key.shape + H_LOCAL = H // world_size + query = query.reshape(B, S_Q_LOCAL, world_size, H_LOCAL, D).permute(2, 1, 0, 3, 4).contiguous() + key = key.reshape(B, S_KV_LOCAL, world_size, H_LOCAL, D).permute(2, 1, 0, 3, 4).contiguous() + value = value.reshape(B, S_KV_LOCAL, world_size, H_LOCAL, D).permute(2, 1, 0, 3, 4).contiguous() + query, key, value = (_all_to_all_single(x, group) for x in (query, key, value)) + query, key, value = (x.flatten(0, 1).permute(1, 0, 2, 3).contiguous() for x in (query, key, value)) + + out = forward_op( + ctx, + query, + key, + value, + attn_mask, + dropout_p, + is_causal, + scale, + enable_gqa, + return_lse, + _save_ctx=True, + _parallel_config=_parallel_config, + ) + if return_lse: + out, lse, *_ = out + + out = out.reshape(B, world_size, S_Q_LOCAL, H_LOCAL, D).permute(1, 3, 0, 2, 4).contiguous() + out = _all_to_all_single(out, group) + out = out.flatten(0, 1).permute(1, 2, 0, 3).contiguous() + + if return_lse: + lse = lse.reshape(B, world_size, S_Q_LOCAL, H_LOCAL).permute(1, 3, 0, 2).contiguous() + lse = _all_to_all_single(lse, group) + lse = lse.flatten(0, 1).permute(1, 2, 0).contiguous() + else: + lse = None + + return (out, lse) if return_lse else out + + @staticmethod + def backward( + ctx: torch.autograd.function.FunctionCtx, + grad_out: torch.Tensor, + *args, + ): + ulysses_mesh = ctx._parallel_config.context_parallel_config._ulysses_mesh + world_size = ctx._parallel_config.context_parallel_config.ulysses_degree + group = ulysses_mesh.get_group() + + B, S_LOCAL, H, D = grad_out.shape + H_LOCAL = H // world_size + + grad_out = grad_out.reshape(B, S_LOCAL, world_size, H_LOCAL, D).permute(2, 1, 0, 3, 4).contiguous() + grad_out = _all_to_all_single(grad_out, group) + grad_out = grad_out.flatten(0, 1).permute(1, 0, 2, 3).contiguous() + + grad_query_op, grad_key_op, grad_value_op, *_ = ctx.backward_op(ctx, grad_out) + + grad_query, grad_key, grad_value = ( + x.reshape(B, world_size, S_LOCAL, H_LOCAL, D).permute(1, 3, 0, 2, 4).contiguous() + for x in (grad_query_op, grad_key_op, grad_value_op) + ) + grad_query, grad_key, grad_value = (_all_to_all_single(x, group) for x in (grad_query, grad_key, grad_value)) + grad_query, grad_key, grad_value = ( + x.flatten(0, 1).permute(1, 2, 0, 3).contiguous() for x in (grad_query, grad_key, grad_value) + ) + + return grad_query, grad_key, grad_value, None, None, None, None, None, None, None, None + + +def _templated_context_parallel_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: Optional[torch.Tensor] = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: Optional[float] = None, + enable_gqa: bool = False, + return_lse: bool = False, + *, + forward_op, + backward_op, + _parallel_config: Optional["ParallelConfig"] = None, +): + if attn_mask is not None: + raise ValueError("Attention mask is not yet supported for templated attention.") + if is_causal: + raise ValueError("Causal attention is not yet supported for templated attention.") + if enable_gqa: + raise ValueError("GQA is not yet supported for templated attention.") + + # TODO: add support for unified attention with ring/ulysses degree both being > 1 + if _parallel_config.context_parallel_config.ring_degree > 1: + return TemplatedRingAttention.apply( + query, + key, + value, + attn_mask, + dropout_p, + is_causal, + scale, + enable_gqa, + return_lse, + forward_op, + backward_op, + _parallel_config, + ) + elif _parallel_config.context_parallel_config.ulysses_degree > 1: + return TemplatedUlyssesAttention.apply( + query, + key, + value, + attn_mask, + dropout_p, + is_causal, + scale, + enable_gqa, + return_lse, + forward_op, + backward_op, + _parallel_config, + ) + else: + raise ValueError("Reaching this branch of code is unexpected. Please report a bug.") # ===== Attention backends ===== @@ -545,34 +1179,50 @@ def _(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> Tuple[torc @_AttentionBackendRegistry.register( AttentionBackendName.FLASH, constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape], + supports_context_parallel=True, ) def _flash_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, dropout_p: float = 0.0, - scale: Optional[float] = None, is_causal: bool = False, - window_size: Tuple[int, int] = (-1, -1), - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - deterministic: bool = False, - return_attn_probs: bool = False, + scale: Optional[float] = None, + return_lse: bool = False, + _parallel_config: Optional["ParallelConfig"] = None, ) -> torch.Tensor: - out = flash_attn_func( - q=query, - k=key, - v=value, - dropout_p=dropout_p, - softmax_scale=scale, - causal=is_causal, - window_size=window_size, - softcap=softcap, - alibi_slopes=alibi_slopes, - deterministic=deterministic, - return_attn_probs=return_attn_probs, - ) - return out + lse = None + if _parallel_config is None: + out = flash_attn_func( + q=query, + k=key, + v=value, + dropout_p=dropout_p, + softmax_scale=scale, + causal=is_causal, + return_attn_probs=return_lse, + ) + if return_lse: + out, lse, *_ = out + else: + out = _templated_context_parallel_attention( + query, + key, + value, + None, + dropout_p, + is_causal, + scale, + False, + return_lse, + forward_op=_flash_attention_forward_op, + backward_op=_flash_attention_backward_op, + _parallel_config=_parallel_config, + ) + if return_lse: + out, lse = out + + return (out, lse) if return_lse else out @_AttentionBackendRegistry.register( @@ -583,19 +1233,12 @@ def _flash_varlen_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - cu_seqlens_q: Optional[torch.Tensor] = None, - cu_seqlens_k: Optional[torch.Tensor] = None, - max_seqlen_q: Optional[int] = None, - max_seqlen_k: Optional[int] = None, + attn_mask: Optional[torch.Tensor] = None, dropout_p: float = 0.0, scale: Optional[float] = None, is_causal: bool = False, - window_size: Tuple[int, int] = (-1, -1), - softcap: float = 0.0, - alibi_slopes: Optional[torch.Tensor] = None, - deterministic: bool = False, - return_attn_probs: bool = False, - attn_mask: Optional[torch.Tensor] = None, + return_lse: bool = False, + _parallel_config: Optional["ParallelConfig"] = None, ) -> torch.Tensor: batch_size, seq_len_q, _, _ = query.shape _, seq_len_kv, _, _ = key.shape @@ -603,16 +1246,11 @@ def _flash_varlen_attention( if attn_mask is not None: attn_mask = _normalize_attn_mask(attn_mask, batch_size, seq_len_kv) - if any(x is None for x in (cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k)): - (_, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = ( - _prepare_for_flash_attn_or_sage_varlen( - batch_size, seq_len_q, seq_len_kv, attn_mask=attn_mask, device=query.device - ) + (_, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = ( + _prepare_for_flash_attn_or_sage_varlen( + batch_size, seq_len_q, seq_len_kv, attn_mask=attn_mask, device=query.device ) - else: - seqlens_k = torch.full((batch_size,), max_seqlen_k, dtype=torch.int32, device=query.device) - cu_seqlens_q = cu_seqlens_q.to(dtype=torch.int32, device=query.device) - cu_seqlens_k = cu_seqlens_k.to(dtype=torch.int32, device=query.device) + ) key_valid, value_valid = [], [] for b in range(batch_size): @@ -635,11 +1273,7 @@ def _flash_varlen_attention( dropout_p=dropout_p, softmax_scale=scale, causal=is_causal, - window_size=window_size, - softcap=softcap, - alibi_slopes=alibi_slopes, - deterministic=deterministic, - return_attn_probs=return_attn_probs, + return_attn_probs=return_lse, ) out = out.unflatten(0, (batch_size, -1)) @@ -656,30 +1290,17 @@ def _flash_attention_3( value: torch.Tensor, scale: Optional[float] = None, is_causal: bool = False, - window_size: Tuple[int, int] = (-1, -1), - softcap: float = 0.0, - deterministic: bool = False, - return_attn_probs: bool = False, + return_lse: bool = False, + _parallel_config: Optional["ParallelConfig"] = None, ) -> torch.Tensor: - out, lse, *_ = flash_attn_3_func( + out, lse = _wrapped_flash_attn_3( q=query, k=key, v=value, softmax_scale=scale, causal=is_causal, - qv=None, - q_descale=None, - k_descale=None, - v_descale=None, - window_size=window_size, - attention_chunk=0, - softcap=softcap, - num_splits=1, - pack_gqa=None, - deterministic=deterministic, - sm_margin=0, ) - return (out, lse) if return_attn_probs else out + return (out, lse) if return_lse else out @_AttentionBackendRegistry.register( @@ -696,6 +1317,7 @@ def _flash_attention_3_hub( softcap: float = 0.0, deterministic: bool = False, return_attn_probs: bool = False, + _parallel_config: Optional["ParallelConfig"] = None, ) -> torch.Tensor: out = flash_attn_3_func_hub( q=query, @@ -728,17 +1350,11 @@ def _flash_varlen_attention_3( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - cu_seqlens_q: Optional[torch.Tensor] = None, - cu_seqlens_k: Optional[torch.Tensor] = None, - max_seqlen_q: Optional[int] = None, - max_seqlen_k: Optional[int] = None, + attn_mask: Optional[torch.Tensor] = None, scale: Optional[float] = None, is_causal: bool = False, - window_size: Tuple[int, int] = (-1, -1), - softcap: float = 0.0, - deterministic: bool = False, - return_attn_probs: bool = False, - attn_mask: Optional[torch.Tensor] = None, + return_lse: bool = False, + _parallel_config: Optional["ParallelConfig"] = None, ) -> torch.Tensor: batch_size, seq_len_q, _, _ = query.shape _, seq_len_kv, _, _ = key.shape @@ -746,16 +1362,11 @@ def _flash_varlen_attention_3( if attn_mask is not None: attn_mask = _normalize_attn_mask(attn_mask, batch_size, seq_len_kv) - if any(x is None for x in (cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k)): - (_, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = ( - _prepare_for_flash_attn_or_sage_varlen( - batch_size, seq_len_q, seq_len_kv, attn_mask=attn_mask, device=query.device - ) + (_, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = ( + _prepare_for_flash_attn_or_sage_varlen( + batch_size, seq_len_q, seq_len_kv, attn_mask=attn_mask, device=query.device ) - else: - seqlens_k = torch.full((batch_size,), max_seqlen_k, dtype=torch.int32, device=query.device) - cu_seqlens_q = cu_seqlens_q.to(dtype=torch.int32, device=query.device) - cu_seqlens_k = cu_seqlens_k.to(dtype=torch.int32, device=query.device) + ) key_valid, value_valid = [], [] for b in range(batch_size): @@ -775,24 +1386,12 @@ def _flash_varlen_attention_3( cu_seqlens_k=cu_seqlens_k, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, - seqused_q=None, - seqused_k=None, softmax_scale=scale, causal=is_causal, - qv=None, - q_descale=None, - k_descale=None, - v_descale=None, - window_size=window_size, - softcap=softcap, - num_splits=1, - pack_gqa=None, - deterministic=deterministic, - sm_margin=0, ) out = out.unflatten(0, (batch_size, -1)) - return (out, lse) if return_attn_probs else out + return (out, lse) if return_lse else out @_AttentionBackendRegistry.register( @@ -808,7 +1407,7 @@ def _native_flex_attention( scale: Optional[float] = None, enable_gqa: bool = False, return_lse: bool = False, - kernel_options: Optional[Dict[str, Any]] = None, + _parallel_config: Optional["ParallelConfig"] = None, ) -> torch.Tensor: # TODO: should we LRU cache the block mask creation? score_mod = None @@ -853,7 +1452,6 @@ def _native_flex_attention( scale=scale, enable_gqa=enable_gqa, return_lse=return_lse, - kernel_options=kernel_options, ) out = out.permute(0, 2, 1, 3) return out @@ -872,7 +1470,11 @@ def _native_attention( is_causal: bool = False, scale: Optional[float] = None, enable_gqa: bool = False, + return_lse: bool = False, + _parallel_config: Optional["ParallelConfig"] = None, ) -> torch.Tensor: + if return_lse: + raise ValueError("Native attention backend does not support setting `return_lse=True`.") query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value)) out = torch.nn.functional.scaled_dot_product_attention( query=query, @@ -891,6 +1493,7 @@ def _native_attention( @_AttentionBackendRegistry.register( AttentionBackendName._NATIVE_CUDNN, constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape], + supports_context_parallel=True, ) def _native_cudnn_attention( query: torch.Tensor, @@ -901,21 +1504,43 @@ def _native_cudnn_attention( is_causal: bool = False, scale: Optional[float] = None, enable_gqa: bool = False, + return_lse: bool = False, + _parallel_config: Optional["ParallelConfig"] = None, ) -> torch.Tensor: - query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value)) - with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.CUDNN_ATTENTION): - out = torch.nn.functional.scaled_dot_product_attention( - query=query, - key=key, - value=value, - attn_mask=attn_mask, - dropout_p=dropout_p, - is_causal=is_causal, - scale=scale, - enable_gqa=enable_gqa, + lse = None + if _parallel_config is None and not return_lse: + query, key, value = (x.permute(0, 2, 1, 3).contiguous() for x in (query, key, value)) + with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.CUDNN_ATTENTION): + out = torch.nn.functional.scaled_dot_product_attention( + query=query, + key=key, + value=value, + attn_mask=attn_mask, + dropout_p=dropout_p, + is_causal=is_causal, + scale=scale, + enable_gqa=enable_gqa, + ) + out = out.permute(0, 2, 1, 3) + else: + out = _templated_context_parallel_attention( + query, + key, + value, + attn_mask, + dropout_p, + is_causal, + scale, + enable_gqa, + return_lse, + forward_op=_cudnn_attention_forward_op, + backward_op=_cudnn_attention_backward_op, + _parallel_config=_parallel_config, ) - out = out.permute(0, 2, 1, 3) - return out + if return_lse: + out, lse = out + + return (out, lse) if return_lse else out @_AttentionBackendRegistry.register( @@ -931,7 +1556,11 @@ def _native_efficient_attention( is_causal: bool = False, scale: Optional[float] = None, enable_gqa: bool = False, + return_lse: bool = False, + _parallel_config: Optional["ParallelConfig"] = None, ) -> torch.Tensor: + if return_lse: + raise ValueError("Native efficient attention backend does not support setting `return_lse=True`.") query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value)) with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION): out = torch.nn.functional.scaled_dot_product_attention( @@ -960,7 +1589,11 @@ def _native_flash_attention( is_causal: bool = False, scale: Optional[float] = None, enable_gqa: bool = False, + return_lse: bool = False, + _parallel_config: Optional["ParallelConfig"] = None, ) -> torch.Tensor: + if return_lse: + raise ValueError("Native flash attention backend does not support setting `return_lse=True`.") query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value)) with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.FLASH_ATTENTION): out = torch.nn.functional.scaled_dot_product_attention( @@ -990,7 +1623,11 @@ def _native_math_attention( is_causal: bool = False, scale: Optional[float] = None, enable_gqa: bool = False, + return_lse: bool = False, + _parallel_config: Optional["ParallelConfig"] = None, ) -> torch.Tensor: + if return_lse: + raise ValueError("Native math attention backend does not support setting `return_lse=True`.") query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value)) with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH): out = torch.nn.functional.scaled_dot_product_attention( @@ -1017,7 +1654,11 @@ def _native_npu_attention( value: torch.Tensor, dropout_p: float = 0.0, scale: Optional[float] = None, + return_lse: bool = False, + _parallel_config: Optional["ParallelConfig"] = None, ) -> torch.Tensor: + if return_lse: + raise ValueError("NPU attention backend does not support setting `return_lse=True`.") query, key, value = (x.transpose(1, 2).contiguous() for x in (query, key, value)) out = npu_fusion_attention( query, @@ -1047,7 +1688,11 @@ def _native_xla_attention( key: torch.Tensor, value: torch.Tensor, is_causal: bool = False, + return_lse: bool = False, + _parallel_config: Optional["ParallelConfig"] = None, ) -> torch.Tensor: + if return_lse: + raise ValueError("XLA attention backend does not support setting `return_lse=True`.") query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value)) query = query / math.sqrt(query.shape[-1]) out = xla_flash_attention( @@ -1063,6 +1708,7 @@ def _native_xla_attention( @_AttentionBackendRegistry.register( AttentionBackendName.SAGE, constraints=[_check_device_cuda, _check_qkv_dtype_bf16_or_fp16, _check_shape], + supports_context_parallel=True, ) def _sage_attention( query: torch.Tensor, @@ -1071,16 +1717,40 @@ def _sage_attention( is_causal: bool = False, scale: Optional[float] = None, return_lse: bool = False, + _parallel_config: Optional["ParallelConfig"] = None, ) -> torch.Tensor: - return sageattn( - q=query, - k=key, - v=value, - tensor_layout="NHD", - is_causal=is_causal, - sm_scale=scale, - return_lse=return_lse, - ) + lse = None + if _parallel_config is None: + out = sageattn( + q=query, + k=key, + v=value, + tensor_layout="NHD", + is_causal=is_causal, + sm_scale=scale, + return_lse=return_lse, + ) + if return_lse: + out, lse, *_ = out + else: + out = _templated_context_parallel_attention( + query, + key, + value, + None, + 0.0, + is_causal, + scale, + False, + return_lse, + forward_op=_sage_attention_forward_op, + backward_op=_sage_attention_backward_op, + _parallel_config=_parallel_config, + ) + if return_lse: + out, lse = out + + return (out, lse) if return_lse else out @_AttentionBackendRegistry.register( @@ -1091,31 +1761,26 @@ def _sage_varlen_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - cu_seqlens_q: Optional[torch.Tensor] = None, - cu_seqlens_k: Optional[torch.Tensor] = None, - max_seqlen_q: Optional[int] = None, - max_seqlen_k: Optional[int] = None, + attn_mask: Optional[torch.Tensor] = None, is_causal: bool = False, scale: Optional[float] = None, - smooth_k: bool = True, - attn_mask: Optional[torch.Tensor] = None, + return_lse: bool = False, + _parallel_config: Optional["ParallelConfig"] = None, ) -> torch.Tensor: + if return_lse: + raise ValueError("Sage varlen backend does not support setting `return_lse=True`.") + batch_size, seq_len_q, _, _ = query.shape _, seq_len_kv, _, _ = key.shape if attn_mask is not None: attn_mask = _normalize_attn_mask(attn_mask, batch_size, seq_len_kv) - if any(x is None for x in (cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k)): - (_, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = ( - _prepare_for_flash_attn_or_sage_varlen( - batch_size, seq_len_q, seq_len_kv, attn_mask=attn_mask, device=query.device - ) + (_, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = ( + _prepare_for_flash_attn_or_sage_varlen( + batch_size, seq_len_q, seq_len_kv, attn_mask=attn_mask, device=query.device ) - else: - seqlens_k = torch.full((batch_size,), max_seqlen_k, dtype=torch.int32, device=query.device) - cu_seqlens_q = cu_seqlens_q.to(dtype=torch.int32, device=query.device) - cu_seqlens_k = cu_seqlens_k.to(dtype=torch.int32, device=query.device) + ) key_valid, value_valid = [], [] for b in range(batch_size): @@ -1137,7 +1802,6 @@ def _sage_varlen_attention( max_seqlen_k=max_seqlen_k, is_causal=is_causal, sm_scale=scale, - smooth_k=smooth_k, ) out = out.unflatten(0, (batch_size, -1)) @@ -1154,11 +1818,8 @@ def _sage_qk_int8_pv_fp8_cuda_attention( value: torch.Tensor, is_causal: bool = False, scale: Optional[float] = None, - qk_quant_gran: _SAGE_ATTENTION_QK_QUANT_GRAN = "per_thread", - pv_accum_dtype: _SAGE_ATTENTION_PV_ACCUM_DTYPE = "fp32+fp32", - smooth_k: bool = True, - smooth_v: bool = False, return_lse: bool = False, + _parallel_config: Optional["ParallelConfig"] = None, ) -> torch.Tensor: return sageattn_qk_int8_pv_fp8_cuda( q=query, @@ -1166,11 +1827,7 @@ def _sage_qk_int8_pv_fp8_cuda_attention( v=value, tensor_layout="NHD", is_causal=is_causal, - qk_quant_gran=qk_quant_gran, sm_scale=scale, - pv_accum_dtype=pv_accum_dtype, - smooth_k=smooth_k, - smooth_v=smooth_v, return_lse=return_lse, ) @@ -1185,10 +1842,8 @@ def _sage_qk_int8_pv_fp8_cuda_sm90_attention( value: torch.Tensor, is_causal: bool = False, scale: Optional[float] = None, - qk_quant_gran: _SAGE_ATTENTION_QK_QUANT_GRAN = "per_thread", - pv_accum_dtype: _SAGE_ATTENTION_PV_ACCUM_DTYPE = "fp32+fp32", - smooth_k: bool = True, return_lse: bool = False, + _parallel_config: Optional["ParallelConfig"] = None, ) -> torch.Tensor: return sageattn_qk_int8_pv_fp8_cuda_sm90( q=query, @@ -1196,10 +1851,7 @@ def _sage_qk_int8_pv_fp8_cuda_sm90_attention( v=value, tensor_layout="NHD", is_causal=is_causal, - qk_quant_gran=qk_quant_gran, sm_scale=scale, - pv_accum_dtype=pv_accum_dtype, - smooth_k=smooth_k, return_lse=return_lse, ) @@ -1214,11 +1866,8 @@ def _sage_qk_int8_pv_fp16_cuda_attention( value: torch.Tensor, is_causal: bool = False, scale: Optional[float] = None, - qk_quant_gran: _SAGE_ATTENTION_QK_QUANT_GRAN = "per_thread", - pv_accum_dtype: _SAGE_ATTENTION_PV_ACCUM_DTYPE = "fp32", - smooth_k: bool = True, - smooth_v: bool = False, return_lse: bool = False, + _parallel_config: Optional["ParallelConfig"] = None, ) -> torch.Tensor: return sageattn_qk_int8_pv_fp16_cuda( q=query, @@ -1226,11 +1875,7 @@ def _sage_qk_int8_pv_fp16_cuda_attention( v=value, tensor_layout="NHD", is_causal=is_causal, - qk_quant_gran=qk_quant_gran, sm_scale=scale, - pv_accum_dtype=pv_accum_dtype, - smooth_k=smooth_k, - smooth_v=smooth_v, return_lse=return_lse, ) @@ -1245,19 +1890,16 @@ def _sage_qk_int8_pv_fp16_triton_attention( value: torch.Tensor, is_causal: bool = False, scale: Optional[float] = None, - quantization_backend: _SAGE_ATTENTION_QUANTIZATION_BACKEND = "triton", - smooth_k: bool = True, return_lse: bool = False, + _parallel_config: Optional["ParallelConfig"] = None, ) -> torch.Tensor: return sageattn_qk_int8_pv_fp16_triton( q=query, k=key, v=value, tensor_layout="NHD", - quantization_backend=quantization_backend, is_causal=is_causal, sm_scale=scale, - smooth_k=smooth_k, return_lse=return_lse, ) @@ -1275,7 +1917,12 @@ def _xformers_attention( is_causal: bool = False, scale: Optional[float] = None, enable_gqa: bool = False, + return_lse: bool = False, + _parallel_config: Optional["ParallelConfig"] = None, ) -> torch.Tensor: + if return_lse: + raise ValueError("xformers attention backend does not support setting `return_lse=True`.") + batch_size, seq_len_q, num_heads_q, _ = query.shape _, seq_len_kv, num_heads_kv, _ = key.shape diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index 2388989be2..b3d74954bd 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -65,6 +65,7 @@ from ..utils.hub_utils import ( populate_model_card, ) from ..utils.torch_utils import empty_device_cache +from ._modeling_parallel import ContextParallelConfig, ContextParallelModelPlan, ParallelConfig from .model_loading_utils import ( _caching_allocator_warmup, _determine_device_map, @@ -248,6 +249,8 @@ class ModelMixin(torch.nn.Module, PushToHubMixin): _skip_layerwise_casting_patterns = None _supports_group_offloading = True _repeated_blocks = [] + _parallel_config = None + _cp_plan = None def __init__(self): super().__init__() @@ -620,8 +623,8 @@ class ModelMixin(torch.nn.Module, PushToHubMixin): def reset_attention_backend(self) -> None: """ - Resets the attention backend for the model. Following calls to `forward` will use the environment default or - the torch native scaled dot product attention. + Resets the attention backend for the model. Following calls to `forward` will use the environment default, if + set, or the torch native scaled dot product attention. """ from .attention import AttentionModuleMixin from .attention_processor import Attention, MochiAttention @@ -960,6 +963,7 @@ class ModelMixin(torch.nn.Module, PushToHubMixin): quantization_config = kwargs.pop("quantization_config", None) dduf_entries: Optional[Dict[str, DDUFEntry]] = kwargs.pop("dduf_entries", None) disable_mmap = kwargs.pop("disable_mmap", False) + parallel_config: Optional[Union[ParallelConfig, ContextParallelConfig]] = kwargs.pop("parallel_config", None) is_parallel_loading_enabled = HF_ENABLE_PARALLEL_LOADING if is_parallel_loading_enabled and not low_cpu_mem_usage: @@ -1340,6 +1344,9 @@ class ModelMixin(torch.nn.Module, PushToHubMixin): # Set model in evaluation mode to deactivate DropOut modules by default model.eval() + if parallel_config is not None: + model.enable_parallelism(config=parallel_config) + if output_loading_info: return model, loading_info @@ -1478,6 +1485,73 @@ class ModelMixin(torch.nn.Module, PushToHubMixin): f"Regional compilation failed because {repeated_blocks} classes are not found in the model. " ) + def enable_parallelism( + self, + *, + config: Union[ParallelConfig, ContextParallelConfig], + cp_plan: Optional[Dict[str, ContextParallelModelPlan]] = None, + ): + from ..hooks.context_parallel import apply_context_parallel + from .attention import AttentionModuleMixin + from .attention_processor import Attention, MochiAttention + + logger.warning( + "`enable_parallelism` is an experimental feature. The API may change in the future and breaking changes may be introduced at any time without warning." + ) + + if isinstance(config, ContextParallelConfig): + config = ParallelConfig(context_parallel_config=config) + + if not torch.distributed.is_initialized(): + raise RuntimeError("torch.distributed must be initialized before calling `enable_parallelism`.") + + rank = torch.distributed.get_rank() + world_size = torch.distributed.get_world_size() + device_type = torch._C._get_accelerator().type + device_module = torch.get_device_module(device_type) + device = torch.device(device_type, rank % device_module.device_count()) + + cp_mesh = None + if config.context_parallel_config is not None: + cp_config = config.context_parallel_config + if cp_config.ring_degree < 1 or cp_config.ulysses_degree < 1: + raise ValueError("`ring_degree` and `ulysses_degree` must be greater than or equal to 1.") + if cp_config.ring_degree > 1 and cp_config.ulysses_degree > 1: + raise ValueError( + "Unified Ulysses-Ring attention is not yet supported. Please set either `ring_degree` or `ulysses_degree` to 1." + ) + if cp_config.ring_degree * cp_config.ulysses_degree > world_size: + raise ValueError( + f"The product of `ring_degree` ({cp_config.ring_degree}) and `ulysses_degree` ({cp_config.ulysses_degree}) must not exceed the world size ({world_size})." + ) + cp_mesh = torch.distributed.device_mesh.init_device_mesh( + device_type=device_type, + mesh_shape=(cp_config.ring_degree, cp_config.ulysses_degree), + mesh_dim_names=("ring", "ulysses"), + ) + + config.setup(rank, world_size, device, cp_mesh=cp_mesh) + + if cp_plan is None and self._cp_plan is None: + raise ValueError( + "`cp_plan` must be provided either as an argument or set in the model's `_cp_plan` attribute." + ) + cp_plan = cp_plan if cp_plan is not None else self._cp_plan + + if config.context_parallel_config is not None: + apply_context_parallel(self, config.context_parallel_config, cp_plan) + + self._parallel_config = config + + attention_classes = (Attention, MochiAttention, AttentionModuleMixin) + for module in self.modules(): + if not isinstance(module, attention_classes): + continue + processor = module.processor + if processor is None or not hasattr(processor, "_parallel_config"): + continue + processor._parallel_config = config + @classmethod def _load_pretrained_model( cls, diff --git a/src/diffusers/models/transformers/transformer_bria.py b/src/diffusers/models/transformers/transformer_bria.py index 04a9c5645c..d54679306e 100644 --- a/src/diffusers/models/transformers/transformer_bria.py +++ b/src/diffusers/models/transformers/transformer_bria.py @@ -120,6 +120,7 @@ def get_1d_rotary_pos_embed( class BriaAttnProcessor: _attention_backend = None + _parallel_config = None def __init__(self): if not hasattr(F, "scaled_dot_product_attention"): @@ -161,7 +162,12 @@ class BriaAttnProcessor: key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1) hidden_states = dispatch_attention_fn( - query, key, value, attn_mask=attention_mask, backend=self._attention_backend + query, + key, + value, + attn_mask=attention_mask, + backend=self._attention_backend, + parallel_config=self._parallel_config, ) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(query.dtype) diff --git a/src/diffusers/models/transformers/transformer_flux.py b/src/diffusers/models/transformers/transformer_flux.py index 7ab371a1a1..1a44644324 100644 --- a/src/diffusers/models/transformers/transformer_flux.py +++ b/src/diffusers/models/transformers/transformer_flux.py @@ -24,6 +24,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FluxTransformer2DLoadersMixin, FromOriginalModelMixin, PeftAdapterMixin from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers from ...utils.torch_utils import maybe_allow_in_graph +from .._modeling_parallel import ContextParallelInput, ContextParallelOutput from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward from ..attention_dispatch import dispatch_attention_fn from ..cache_utils import CacheMixin @@ -73,6 +74,7 @@ def _get_qkv_projections(attn: "FluxAttention", hidden_states, encoder_hidden_st class FluxAttnProcessor: _attention_backend = None + _parallel_config = None def __init__(self): if not hasattr(F, "scaled_dot_product_attention"): @@ -114,7 +116,12 @@ class FluxAttnProcessor: key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1) hidden_states = dispatch_attention_fn( - query, key, value, attn_mask=attention_mask, backend=self._attention_backend + query, + key, + value, + attn_mask=attention_mask, + backend=self._attention_backend, + parallel_config=self._parallel_config, ) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(query.dtype) @@ -136,6 +143,7 @@ class FluxIPAdapterAttnProcessor(torch.nn.Module): """Flux Attention processor for IP-Adapter.""" _attention_backend = None + _parallel_config = None def __init__( self, hidden_size: int, cross_attention_dim: int, num_tokens=(4,), scale=1.0, device=None, dtype=None @@ -220,6 +228,7 @@ class FluxIPAdapterAttnProcessor(torch.nn.Module): dropout_p=0.0, is_causal=False, backend=self._attention_backend, + parallel_config=self._parallel_config, ) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(query.dtype) @@ -252,6 +261,7 @@ class FluxIPAdapterAttnProcessor(torch.nn.Module): dropout_p=0.0, is_causal=False, backend=self._attention_backend, + parallel_config=self._parallel_config, ) current_ip_hidden_states = current_ip_hidden_states.reshape(batch_size, -1, attn.heads * attn.head_dim) current_ip_hidden_states = current_ip_hidden_states.to(ip_query.dtype) @@ -556,6 +566,15 @@ class FluxTransformer2DModel( _no_split_modules = ["FluxTransformerBlock", "FluxSingleTransformerBlock"] _skip_layerwise_casting_patterns = ["pos_embed", "norm"] _repeated_blocks = ["FluxTransformerBlock", "FluxSingleTransformerBlock"] + _cp_plan = { + "": { + "hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False), + "encoder_hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False), + "img_ids": ContextParallelInput(split_dim=0, expected_dims=2, split_output=False), + "txt_ids": ContextParallelInput(split_dim=0, expected_dims=2, split_output=False), + }, + "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3), + } @register_to_config def __init__( diff --git a/src/diffusers/models/transformers/transformer_ltx.py b/src/diffusers/models/transformers/transformer_ltx.py index 79149fb760..9f3840690d 100644 --- a/src/diffusers/models/transformers/transformer_ltx.py +++ b/src/diffusers/models/transformers/transformer_ltx.py @@ -24,6 +24,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ...utils import USE_PEFT_BACKEND, deprecate, is_torch_version, logging, scale_lora_layers, unscale_lora_layers from ...utils.torch_utils import maybe_allow_in_graph +from .._modeling_parallel import ContextParallelInput, ContextParallelOutput from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward from ..attention_dispatch import dispatch_attention_fn from ..cache_utils import CacheMixin @@ -51,6 +52,7 @@ class LTXVideoAttnProcessor: """ _attention_backend = None + _parallel_config = None def __init__(self): if is_torch_version("<", "2.0"): @@ -100,6 +102,7 @@ class LTXVideoAttnProcessor: dropout_p=0.0, is_causal=False, backend=self._attention_backend, + parallel_config=self._parallel_config, ) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(query.dtype) @@ -409,6 +412,18 @@ class LTXVideoTransformer3DModel( _supports_gradient_checkpointing = True _skip_layerwise_casting_patterns = ["norm"] _repeated_blocks = ["LTXVideoTransformerBlock"] + _cp_plan = { + "": { + "hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False), + "encoder_hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False), + "encoder_attention_mask": ContextParallelInput(split_dim=1, expected_dims=2, split_output=False), + }, + "rope": { + 0: ContextParallelInput(split_dim=1, expected_dims=3, split_output=True), + 1: ContextParallelInput(split_dim=1, expected_dims=3, split_output=True), + }, + "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3), + } @register_to_config def __init__( diff --git a/src/diffusers/models/transformers/transformer_qwenimage.py b/src/diffusers/models/transformers/transformer_qwenimage.py index 846add8906..05379270c1 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage.py +++ b/src/diffusers/models/transformers/transformer_qwenimage.py @@ -25,6 +25,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers from ...utils.torch_utils import maybe_allow_in_graph +from .._modeling_parallel import ContextParallelInput, ContextParallelOutput from ..attention import AttentionMixin, FeedForward from ..attention_dispatch import dispatch_attention_fn from ..attention_processor import Attention @@ -261,6 +262,7 @@ class QwenDoubleStreamAttnProcessor2_0: """ _attention_backend = None + _parallel_config = None def __init__(self): if not hasattr(F, "scaled_dot_product_attention"): @@ -334,6 +336,7 @@ class QwenDoubleStreamAttnProcessor2_0: dropout_p=0.0, is_causal=False, backend=self._attention_backend, + parallel_config=self._parallel_config, ) # Reshape back @@ -502,6 +505,18 @@ class QwenImageTransformer2DModel( _no_split_modules = ["QwenImageTransformerBlock"] _skip_layerwise_casting_patterns = ["pos_embed", "norm"] _repeated_blocks = ["QwenImageTransformerBlock"] + _cp_plan = { + "": { + "hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False), + "encoder_hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False), + "encoder_hidden_states_mask": ContextParallelInput(split_dim=1, expected_dims=2, split_output=False), + }, + "pos_embed": { + 0: ContextParallelInput(split_dim=0, expected_dims=2, split_output=True), + 1: ContextParallelInput(split_dim=0, expected_dims=2, split_output=True), + }, + "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3), + } @register_to_config def __init__( diff --git a/src/diffusers/models/transformers/transformer_skyreels_v2.py b/src/diffusers/models/transformers/transformer_skyreels_v2.py index 358759164b..6b600aa224 100644 --- a/src/diffusers/models/transformers/transformer_skyreels_v2.py +++ b/src/diffusers/models/transformers/transformer_skyreels_v2.py @@ -73,6 +73,7 @@ def _get_added_kv_projections(attn: "SkyReelsV2Attention", encoder_hidden_states class SkyReelsV2AttnProcessor: _attention_backend = None + _parallel_config = None def __init__(self): if not hasattr(F, "scaled_dot_product_attention"): @@ -139,6 +140,7 @@ class SkyReelsV2AttnProcessor: dropout_p=0.0, is_causal=False, backend=self._attention_backend, + parallel_config=self._parallel_config, ) hidden_states_img = hidden_states_img.flatten(2, 3) hidden_states_img = hidden_states_img.type_as(query) @@ -151,6 +153,7 @@ class SkyReelsV2AttnProcessor: dropout_p=0.0, is_causal=False, backend=self._attention_backend, + parallel_config=self._parallel_config, ) hidden_states = hidden_states.flatten(2, 3) diff --git a/src/diffusers/models/transformers/transformer_wan.py b/src/diffusers/models/transformers/transformer_wan.py index 968a0369c2..25c055fb56 100644 --- a/src/diffusers/models/transformers/transformer_wan.py +++ b/src/diffusers/models/transformers/transformer_wan.py @@ -23,6 +23,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ...utils import USE_PEFT_BACKEND, deprecate, logging, scale_lora_layers, unscale_lora_layers from ...utils.torch_utils import maybe_allow_in_graph +from .._modeling_parallel import ContextParallelInput, ContextParallelOutput from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward from ..attention_dispatch import dispatch_attention_fn from ..cache_utils import CacheMixin @@ -66,6 +67,7 @@ def _get_added_kv_projections(attn: "WanAttention", encoder_hidden_states_img: t class WanAttnProcessor: _attention_backend = None + _parallel_config = None def __init__(self): if not hasattr(F, "scaled_dot_product_attention"): @@ -132,6 +134,7 @@ class WanAttnProcessor: dropout_p=0.0, is_causal=False, backend=self._attention_backend, + parallel_config=self._parallel_config, ) hidden_states_img = hidden_states_img.flatten(2, 3) hidden_states_img = hidden_states_img.type_as(query) @@ -144,6 +147,7 @@ class WanAttnProcessor: dropout_p=0.0, is_causal=False, backend=self._attention_backend, + parallel_config=self._parallel_config, ) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.type_as(query) @@ -539,6 +543,19 @@ class WanTransformer3DModel( _keep_in_fp32_modules = ["time_embedder", "scale_shift_table", "norm1", "norm2", "norm3"] _keys_to_ignore_on_load_unexpected = ["norm_added_q"] _repeated_blocks = ["WanTransformerBlock"] + _cp_plan = { + "rope": { + 0: ContextParallelInput(split_dim=1, expected_dims=4, split_output=True), + 1: ContextParallelInput(split_dim=1, expected_dims=4, split_output=True), + }, + "blocks.0": { + "hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False), + }, + "blocks.*": { + "encoder_hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False), + }, + "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3), + } @register_to_config def __init__( diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index bbb9712496..6e7d227979 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -648,6 +648,21 @@ class ConsistencyDecoderVAE(metaclass=DummyObject): requires_backends(cls, ["torch"]) +class ContextParallelConfig(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class ControlNetModel(metaclass=DummyObject): _backends = ["torch"] @@ -1053,6 +1068,21 @@ class OmniGenTransformer2DModel(metaclass=DummyObject): requires_backends(cls, ["torch"]) +class ParallelConfig(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class PixArtTransformer2DModel(metaclass=DummyObject): _backends = ["torch"] From 310fdaf5561d1b20240a2b66e978edb66175ad5c Mon Sep 17 00:00:00 2001 From: DefTruth <31974251+DefTruth@users.noreply.github.com> Date: Thu, 25 Sep 2025 01:50:57 +0800 Subject: [PATCH 026/123] Introduce cache-dit to community optimization (#12366) * docs: introduce cache-dit to diffusers * docs: introduce cache-dit to diffusers * docs: introduce cache-dit to diffusers * docs: introduce cache-dit to diffusers * docs: introduce cache-dit to diffusers * docs: introduce cache-dit to diffusers * docs: introduce cache-dit to diffusers * misc: update examples link * misc: update examples link * docs: introduce cache-dit to diffusers * docs: introduce cache-dit to diffusers * docs: introduce cache-dit to diffusers * docs: introduce cache-dit to diffusers * docs: introduce cache-dit to diffusers * Refine documentation for CacheDiT features Updated the wording for clarity and consistency in the documentation. Adjusted sections on cache acceleration, automatic block adapter, patch functor, and hybrid cache configuration. --- docs/source/en/_toctree.yml | 2 + docs/source/en/optimization/cache_dit.md | 270 +++++++++++++++++++++++ 2 files changed, 272 insertions(+) create mode 100644 docs/source/en/optimization/cache_dit.md diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 07408cc30e..96c6fbb17f 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -82,6 +82,8 @@ title: Token merging - local: optimization/deepcache title: DeepCache + - local: optimization/cache_dit + title: CacheDiT - local: optimization/tgate title: TGATE - local: optimization/xdit diff --git a/docs/source/en/optimization/cache_dit.md b/docs/source/en/optimization/cache_dit.md new file mode 100644 index 0000000000..1261423212 --- /dev/null +++ b/docs/source/en/optimization/cache_dit.md @@ -0,0 +1,270 @@ +## CacheDiT + +CacheDiT is a unified, flexible, and training-free cache acceleration framework designed to support nearly all Diffusers' DiT-based pipelines. It provides a unified cache API that supports automatic block adapter, DBCache, and more. + +To learn more, refer to the [CacheDiT](https://github.com/vipshop/cache-dit) repository. + +Install a stable release of CacheDiT from PyPI or you can install the latest version from GitHub. + + + + +```bash +pip3 install -U cache-dit +``` + + + + +```bash +pip3 install git+https://github.com/vipshop/cache-dit.git +``` + + + + +Run the command below to view supported DiT pipelines. + +```python +>>> import cache_dit +>>> cache_dit.supported_pipelines() +(30, ['Flux*', 'Mochi*', 'CogVideoX*', 'Wan*', 'HunyuanVideo*', 'QwenImage*', 'LTX*', 'Allegro*', +'CogView3Plus*', 'CogView4*', 'Cosmos*', 'EasyAnimate*', 'SkyReelsV2*', 'StableDiffusion3*', +'ConsisID*', 'DiT*', 'Amused*', 'Bria*', 'Lumina*', 'OmniGen*', 'PixArt*', 'Sana*', 'StableAudio*', +'VisualCloze*', 'AuraFlow*', 'Chroma*', 'ShapE*', 'HiDream*', 'HunyuanDiT*', 'HunyuanDiTPAG*']) +``` + +For a complete benchmark, please refer to [Benchmarks](https://github.com/vipshop/cache-dit/blob/main/bench/). + + +## Unified Cache API + +CacheDiT works by matching specific input/output patterns as shown below. + +![](https://github.com/vipshop/cache-dit/raw/main/assets/patterns-v1.png) + +Call the `enable_cache()` function on a pipeline to enable cache acceleration. This function is the entry point to many of CacheDiT's features. + +```python +import cache_dit +from diffusers import DiffusionPipeline + +# Can be any diffusion pipeline +pipe = DiffusionPipeline.from_pretrained("Qwen/Qwen-Image") + +# One-line code with default cache options. +cache_dit.enable_cache(pipe) + +# Just call the pipe as normal. +output = pipe(...) + +# Disable cache and run original pipe. +cache_dit.disable_cache(pipe) +``` + +## Automatic Block Adapter + +For custom or modified pipelines or transformers not included in Diffusers, use the `BlockAdapter` in `auto` mode or via manual configuration. Please check the [BlockAdapter](https://github.com/vipshop/cache-dit/blob/main/docs/User_Guide.md#automatic-block-adapter) docs for more details. Refer to [Qwen-Image w/ BlockAdapter](https://github.com/vipshop/cache-dit/blob/main/examples/adapter/run_qwen_image_adapter.py) as an example. + + +```python +from cache_dit import ForwardPattern, BlockAdapter + +# Use 🔥BlockAdapter with `auto` mode. +cache_dit.enable_cache( + BlockAdapter( + # Any DiffusionPipeline, Qwen-Image, etc. + pipe=pipe, auto=True, + # Check `📚Forward Pattern Matching` documentation and hack the code of + # of Qwen-Image, you will find that it has satisfied `FORWARD_PATTERN_1`. + forward_pattern=ForwardPattern.Pattern_1, + ), +) + +# Or, manually setup transformer configurations. +cache_dit.enable_cache( + BlockAdapter( + pipe=pipe, # Qwen-Image, etc. + transformer=pipe.transformer, + blocks=pipe.transformer.transformer_blocks, + forward_pattern=ForwardPattern.Pattern_1, + ), +) +``` + +Sometimes, a Transformer class will contain more than one transformer `blocks`. For example, FLUX.1 (HiDream, Chroma, etc) contains `transformer_blocks` and `single_transformer_blocks` (with different forward patterns). The BlockAdapter is able to detect this hybrid pattern type as well. +Refer to [FLUX.1](https://github.com/vipshop/cache-dit/blob/main/examples/adapter/run_flux_adapter.py) as an example. + +```python +# For diffusers <= 0.34.0, FLUX.1 transformer_blocks and +# single_transformer_blocks have different forward patterns. +cache_dit.enable_cache( + BlockAdapter( + pipe=pipe, # FLUX.1, etc. + transformer=pipe.transformer, + blocks=[ + pipe.transformer.transformer_blocks, + pipe.transformer.single_transformer_blocks, + ], + forward_pattern=[ + ForwardPattern.Pattern_1, + ForwardPattern.Pattern_3, + ], + ), +) +``` + +This also works if there is more than one transformer (namely `transformer` and `transformer_2`) in its structure. Refer to [Wan 2.2 MoE](https://github.com/vipshop/cache-dit/blob/main/examples/pipeline/run_wan_2.2.py) as an example. + +## Patch Functor + +For any pattern not included in CacheDiT, use the Patch Functor to convert the pattern into a known pattern. You need to subclass the Patch Functor and may also need to fuse the operations within the blocks for loop into block `forward`. After implementing a Patch Functor, set the `patch_functor` property in `BlockAdapter`. + +![](https://github.com/vipshop/cache-dit/raw/main/assets/patch-functor.png) + +Some Patch Functors are already provided in CacheDiT, [HiDreamPatchFunctor](https://github.com/vipshop/cache-dit/blob/main/src/cache_dit/cache_factory/patch_functors/functor_hidream.py), [ChromaPatchFunctor](https://github.com/vipshop/cache-dit/blob/main/src/cache_dit/cache_factory/patch_functors/functor_chroma.py), etc. + +```python +@BlockAdapterRegistry.register("HiDream") +def hidream_adapter(pipe, **kwargs) -> BlockAdapter: + from diffusers import HiDreamImageTransformer2DModel + from cache_dit.cache_factory.patch_functors import HiDreamPatchFunctor + + assert isinstance(pipe.transformer, HiDreamImageTransformer2DModel) + return BlockAdapter( + pipe=pipe, + transformer=pipe.transformer, + blocks=[ + pipe.transformer.double_stream_blocks, + pipe.transformer.single_stream_blocks, + ], + forward_pattern=[ + ForwardPattern.Pattern_0, + ForwardPattern.Pattern_3, + ], + # NOTE: Setup your custom patch functor here. + patch_functor=HiDreamPatchFunctor(), + **kwargs, + ) +``` + +Finally, you can call the `cache_dit.summary()` function on a pipeline after its completed inference to get the cache acceleration details. + +```python +stats = cache_dit.summary(pipe) +``` + +```python +⚡️Cache Steps and Residual Diffs Statistics: QwenImagePipeline + +| Cache Steps | Diffs Min | Diffs P25 | Diffs P50 | Diffs P75 | Diffs P95 | Diffs Max | +|-------------|-----------|-----------|-----------|-----------|-----------|-----------| +| 23 | 0.045 | 0.084 | 0.114 | 0.147 | 0.241 | 0.297 | +``` + +## DBCache: Dual Block Cache + +![](https://github.com/vipshop/cache-dit/raw/main/assets/dbcache-v1.png) + +DBCache (Dual Block Caching) supports different configurations of compute blocks (F8B12, etc.) to enable a balanced trade-off between performance and precision. +- Fn_compute_blocks: Specifies that DBCache uses the **first n** Transformer blocks to fit the information at time step t, enabling the calculation of a more stable L1 diff and delivering more accurate information to subsequent blocks. +- Bn_compute_blocks: Further fuses approximate information in the **last n** Transformer blocks to enhance prediction accuracy. These blocks act as an auto-scaler for approximate hidden states that use residual cache. + + +```python +import cache_dit +from diffusers import FluxPipeline + +pipe_or_adapter = FluxPipeline.from_pretrained( + "black-forest-labs/FLUX.1-dev", + torch_dtype=torch.bfloat16, +).to("cuda") + +# Default options, F8B0, 8 warmup steps, and unlimited cached +# steps for good balance between performance and precision +cache_dit.enable_cache(pipe_or_adapter) + +# Custom options, F8B8, higher precision +from cache_dit import BasicCacheConfig + +cache_dit.enable_cache( + pipe_or_adapter, + cache_config=BasicCacheConfig( + max_warmup_steps=8, # steps do not cache + max_cached_steps=-1, # -1 means no limit + Fn_compute_blocks=8, # Fn, F8, etc. + Bn_compute_blocks=8, # Bn, B8, etc. + residual_diff_threshold=0.12, + ), +) +``` +Check the [DBCache](https://github.com/vipshop/cache-dit/blob/main/docs/DBCache.md) and [User Guide](https://github.com/vipshop/cache-dit/blob/main/docs/User_Guide.md#dbcache) docs for more design details. + +## TaylorSeer Calibrator + +The [TaylorSeers](https://huggingface.co/papers/2503.06923) algorithm further improves the precision of DBCache in cases where the cached steps are large (Hybrid TaylorSeer + DBCache). At timesteps with significant intervals, the feature similarity in diffusion models decreases substantially, significantly harming the generation quality. + +TaylorSeer employs a differential method to approximate the higher-order derivatives of features and predict features in future timesteps with Taylor series expansion. The TaylorSeer implemented in CacheDiT supports both hidden states and residual cache types. F_pred can be a residual cache or a hidden-state cache. + +```python +from cache_dit import BasicCacheConfig, TaylorSeerCalibratorConfig + +cache_dit.enable_cache( + pipe_or_adapter, + # Basic DBCache w/ FnBn configurations + cache_config=BasicCacheConfig( + max_warmup_steps=8, # steps do not cache + max_cached_steps=-1, # -1 means no limit + Fn_compute_blocks=8, # Fn, F8, etc. + Bn_compute_blocks=8, # Bn, B8, etc. + residual_diff_threshold=0.12, + ), + # Then, you can use the TaylorSeer Calibrator to approximate + # the values in cached steps, taylorseer_order default is 1. + calibrator_config=TaylorSeerCalibratorConfig( + taylorseer_order=1, + ), +) +``` + +> [!TIP] +> The `Bn_compute_blocks` parameter of DBCache can be set to `0` if you use TaylorSeer as the calibrator for approximate hidden states. DBCache's `Bn_compute_blocks` also acts as a calibrator, so you can choose either `Bn_compute_blocks` > 0 or TaylorSeer. We recommend using the configuration scheme of TaylorSeer + DBCache FnB0. + +## Hybrid Cache CFG + +CacheDiT supports caching for CFG (classifier-free guidance). For models that fuse CFG and non-CFG into a single forward step, or models that do not include CFG in the forward step, please set `enable_separate_cfg` parameter to `False (default, None)`. Otherwise, set it to `True`. + +```python +from cache_dit import BasicCacheConfig + +cache_dit.enable_cache( + pipe_or_adapter, + cache_config=BasicCacheConfig( + ..., + # For example, set it as True for Wan 2.1, Qwen-Image + # and set it as False for FLUX.1, HunyuanVideo, etc. + enable_separate_cfg=True, + ), +) +``` + +## torch.compile + +CacheDiT is designed to work with torch.compile for even better performance. Call `torch.compile` after enabling the cache. + + +```python +cache_dit.enable_cache(pipe) + +# Compile the Transformer module +pipe.transformer = torch.compile(pipe.transformer) +``` + +If you're using CacheDiT with dynamic input shapes, consider increasing the `recompile_limit` of `torch._dynamo`. Otherwise, the `recompile_limit` error may be triggered, causing the module to fall back to eager mode. + +```python +torch._dynamo.config.recompile_limit = 96 # default is 8 +torch._dynamo.config.accumulated_recompile_limit = 2048 # default is 256 +``` + +Please check [perf.py](https://github.com/vipshop/cache-dit/blob/main/bench/perf.py) for more details. From ec5449f3a1378df207df481bfa1ad7ff8057a58a Mon Sep 17 00:00:00 2001 From: Lucain Date: Thu, 25 Sep 2025 18:28:54 +0200 Subject: [PATCH 027/123] Support both huggingface_hub `v0.x` and `v1.x` (#12389) * Support huggingface_hub 0.x and 1.x * httpx --- setup.py | 4 +++- src/diffusers/configuration_utils.py | 4 ++-- src/diffusers/dependency_versions_table.py | 3 ++- src/diffusers/models/modeling_flax_utils.py | 4 ++-- src/diffusers/pipelines/pipeline_loading_utils.py | 6 +++--- src/diffusers/pipelines/pipeline_utils.py | 6 +++--- src/diffusers/utils/hub_utils.py | 6 +++--- tests/models/test_modeling_common.py | 7 +++---- tests/pipelines/test_pipelines.py | 6 +++--- 9 files changed, 24 insertions(+), 22 deletions(-) diff --git a/setup.py b/setup.py index ba3ad8e2b3..372a568595 100644 --- a/setup.py +++ b/setup.py @@ -102,7 +102,8 @@ _deps = [ "filelock", "flax>=0.4.1", "hf-doc-builder>=0.3.0", - "huggingface-hub>=0.34.0", + "httpx<1.0.0", + "huggingface-hub>=0.34.0,<2.0", "requests-mock==1.10.0", "importlib_metadata", "invisible-watermark>=0.2.0", @@ -259,6 +260,7 @@ extras["dev"] = ( install_requires = [ deps["importlib_metadata"], deps["filelock"], + deps["httpx"], deps["huggingface-hub"], deps["numpy"], deps["regex"], diff --git a/src/diffusers/configuration_utils.py b/src/diffusers/configuration_utils.py index 540aab0307..1c4ee33acb 100644 --- a/src/diffusers/configuration_utils.py +++ b/src/diffusers/configuration_utils.py @@ -30,11 +30,11 @@ import numpy as np from huggingface_hub import DDUFEntry, create_repo, hf_hub_download from huggingface_hub.utils import ( EntryNotFoundError, + HfHubHTTPError, RepositoryNotFoundError, RevisionNotFoundError, validate_hf_hub_args, ) -from requests import HTTPError from typing_extensions import Self from . import __version__ @@ -419,7 +419,7 @@ class ConfigMixin: raise EnvironmentError( f"{pretrained_model_name_or_path} does not appear to have a file named {cls.config_name}." ) - except HTTPError as err: + except HfHubHTTPError as err: raise EnvironmentError( "There was a specific connection error when trying to load" f" {pretrained_model_name_or_path}:\n{err}" diff --git a/src/diffusers/dependency_versions_table.py b/src/diffusers/dependency_versions_table.py index 79dc4c50a0..bfc4e9818b 100644 --- a/src/diffusers/dependency_versions_table.py +++ b/src/diffusers/dependency_versions_table.py @@ -9,7 +9,8 @@ deps = { "filelock": "filelock", "flax": "flax>=0.4.1", "hf-doc-builder": "hf-doc-builder>=0.3.0", - "huggingface-hub": "huggingface-hub>=0.34.0", + "httpx": "httpx<1.0.0", + "huggingface-hub": "huggingface-hub>=0.34.0,<2.0", "requests-mock": "requests-mock==1.10.0", "importlib_metadata": "importlib_metadata", "invisible-watermark": "invisible-watermark>=0.2.0", diff --git a/src/diffusers/models/modeling_flax_utils.py b/src/diffusers/models/modeling_flax_utils.py index 573828dc4b..8050afff27 100644 --- a/src/diffusers/models/modeling_flax_utils.py +++ b/src/diffusers/models/modeling_flax_utils.py @@ -26,11 +26,11 @@ from flax.traverse_util import flatten_dict, unflatten_dict from huggingface_hub import create_repo, hf_hub_download from huggingface_hub.utils import ( EntryNotFoundError, + HfHubHTTPError, RepositoryNotFoundError, RevisionNotFoundError, validate_hf_hub_args, ) -from requests import HTTPError from .. import __version__, is_torch_available from ..utils import ( @@ -385,7 +385,7 @@ class FlaxModelMixin(PushToHubMixin): raise EnvironmentError( f"{pretrained_model_name_or_path} does not appear to have a file named {FLAX_WEIGHTS_NAME}." ) - except HTTPError as err: + except HfHubHTTPError as err: raise EnvironmentError( f"There was a specific connection error when trying to load {pretrained_model_name_or_path}:\n" f"{err}" diff --git a/src/diffusers/pipelines/pipeline_loading_utils.py b/src/diffusers/pipelines/pipeline_loading_utils.py index 388128df0e..b7a3e08105 100644 --- a/src/diffusers/pipelines/pipeline_loading_utils.py +++ b/src/diffusers/pipelines/pipeline_loading_utils.py @@ -19,12 +19,12 @@ import warnings from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Union +import httpx import requests import torch from huggingface_hub import DDUFEntry, ModelCard, model_info, snapshot_download -from huggingface_hub.utils import OfflineModeIsEnabled, validate_hf_hub_args +from huggingface_hub.utils import HfHubHTTPError, OfflineModeIsEnabled, validate_hf_hub_args from packaging import version -from requests.exceptions import HTTPError from .. import __version__ from ..utils import ( @@ -1110,7 +1110,7 @@ def _download_dduf_file( if not local_files_only: try: info = model_info(pretrained_model_name, token=token, revision=revision) - except (HTTPError, OfflineModeIsEnabled, requests.ConnectionError) as e: + except (HfHubHTTPError, OfflineModeIsEnabled, requests.ConnectionError, httpx.NetworkError) as e: logger.warning(f"Couldn't connect to the Hub: {e}.\nWill try to load from local cache.") local_files_only = True model_info_call_error = e # save error to reraise it if model is not cached locally diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index 01b3c56777..3f6e53099b 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -23,6 +23,7 @@ from dataclasses import dataclass from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Union, get_args, get_origin +import httpx import numpy as np import PIL.Image import requests @@ -36,9 +37,8 @@ from huggingface_hub import ( read_dduf_file, snapshot_download, ) -from huggingface_hub.utils import OfflineModeIsEnabled, validate_hf_hub_args +from huggingface_hub.utils import HfHubHTTPError, OfflineModeIsEnabled, validate_hf_hub_args from packaging import version -from requests.exceptions import HTTPError from tqdm.auto import tqdm from typing_extensions import Self @@ -1616,7 +1616,7 @@ class DiffusionPipeline(ConfigMixin, PushToHubMixin): if not local_files_only: try: info = model_info(pretrained_model_name, token=token, revision=revision) - except (HTTPError, OfflineModeIsEnabled, requests.ConnectionError) as e: + except (HfHubHTTPError, OfflineModeIsEnabled, requests.ConnectionError, httpx.NetworkError) as e: logger.warning(f"Couldn't connect to the Hub: {e}.\nWill try to load from local cache.") local_files_only = True model_info_call_error = e # save error to reraise it if model is not cached locally diff --git a/src/diffusers/utils/hub_utils.py b/src/diffusers/utils/hub_utils.py index fcdf49156a..b6e99452aa 100644 --- a/src/diffusers/utils/hub_utils.py +++ b/src/diffusers/utils/hub_utils.py @@ -38,13 +38,13 @@ from huggingface_hub.constants import HF_HUB_DISABLE_TELEMETRY, HF_HUB_OFFLINE from huggingface_hub.file_download import REGEX_COMMIT_HASH from huggingface_hub.utils import ( EntryNotFoundError, + HfHubHTTPError, RepositoryNotFoundError, RevisionNotFoundError, is_jinja_available, validate_hf_hub_args, ) from packaging import version -from requests import HTTPError from .. import __version__ from .constants import ( @@ -316,7 +316,7 @@ def _get_model_file( raise EnvironmentError( f"{pretrained_model_name_or_path} does not appear to have a file named {weights_name}." ) from e - except HTTPError as e: + except HfHubHTTPError as e: raise EnvironmentError( f"There was a specific connection error when trying to load {pretrained_model_name_or_path}:\n{e}" ) from e @@ -432,7 +432,7 @@ def _get_checkpoint_shard_files( # We have already dealt with RepositoryNotFoundError and RevisionNotFoundError when getting the index, so # we don't have to catch them here. We have also dealt with EntryNotFoundError. - except HTTPError as e: + except HfHubHTTPError as e: raise EnvironmentError( f"We couldn't connect to '{HUGGINGFACE_CO_RESOLVE_ENDPOINT}' to load {pretrained_model_name_or_path}. You should try" " again after checking your internet connection." diff --git a/tests/models/test_modeling_common.py b/tests/models/test_modeling_common.py index 5e7be62342..3a008edfe1 100644 --- a/tests/models/test_modeling_common.py +++ b/tests/models/test_modeling_common.py @@ -37,9 +37,8 @@ import torch import torch.nn as nn from accelerate.utils.modeling import _get_proper_dtype, compute_module_sizes, dtype_byte_size from huggingface_hub import ModelCard, delete_repo, snapshot_download, try_to_load_from_cache -from huggingface_hub.utils import is_jinja_available +from huggingface_hub.utils import HfHubHTTPError, is_jinja_available from parameterized import parameterized -from requests.exceptions import HTTPError from diffusers.models import FluxTransformer2DModel, SD3Transformer2DModel, UNet2DConditionModel from diffusers.models.attention_processor import ( @@ -272,7 +271,7 @@ class ModelUtilsTest(unittest.TestCase): response_mock = mock.Mock() response_mock.status_code = 500 response_mock.headers = {} - response_mock.raise_for_status.side_effect = HTTPError + response_mock.raise_for_status.side_effect = HfHubHTTPError("Server down", response=mock.Mock()) response_mock.json.return_value = {} # Download this model to make sure it's in the cache. @@ -296,7 +295,7 @@ class ModelUtilsTest(unittest.TestCase): error_response = mock.Mock( status_code=500, headers={}, - raise_for_status=mock.Mock(side_effect=HTTPError), + raise_for_status=mock.Mock(side_effect=HfHubHTTPError("Server down", response=mock.Mock())), json=mock.Mock(return_value={}), ) diff --git a/tests/pipelines/test_pipelines.py b/tests/pipelines/test_pipelines.py index 09df140f1a..3a69813612 100644 --- a/tests/pipelines/test_pipelines.py +++ b/tests/pipelines/test_pipelines.py @@ -33,9 +33,9 @@ import safetensors.torch import torch import torch.nn as nn from huggingface_hub import snapshot_download +from huggingface_hub.utils import HfHubHTTPError from parameterized import parameterized from PIL import Image -from requests.exceptions import HTTPError from transformers import CLIPImageProcessor, CLIPModel, CLIPTextConfig, CLIPTextModel, CLIPTokenizer from diffusers import ( @@ -430,7 +430,7 @@ class DownloadTests(unittest.TestCase): response_mock = mock.Mock() response_mock.status_code = 500 response_mock.headers = {} - response_mock.raise_for_status.side_effect = HTTPError + response_mock.raise_for_status.side_effect = HfHubHTTPError("Server down", response=mock.Mock()) response_mock.json.return_value = {} # Download this model to make sure it's in the cache. @@ -457,7 +457,7 @@ class DownloadTests(unittest.TestCase): response_mock = mock.Mock() response_mock.status_code = 500 response_mock.headers = {} - response_mock.raise_for_status.side_effect = HTTPError + response_mock.raise_for_status.side_effect = HfHubHTTPError("Server down", response=mock.Mock()) response_mock.json.return_value = {} # first check that with local files only the pipeline can only be used if cached From 4588bbeb4229fd307119257e273a424b370573b1 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Fri, 26 Sep 2025 18:41:17 +0530 Subject: [PATCH 028/123] [CI] disable installing transformers from main in ci for now. (#12397) * disable installing transformers from main in ci for now. * up * u[p --- .github/workflows/pr_modular_tests.yml | 5 +++-- .github/workflows/pr_tests.yml | 12 +++++++----- .github/workflows/pr_tests_gpu.yml | 13 ++++++++----- tests/pipelines/kandinsky/test_kandinsky.py | 4 +++- .../pipelines/kandinsky/test_kandinsky_combined.py | 12 +++++++++--- tests/pipelines/kandinsky/test_kandinsky_img2img.py | 4 +++- tests/pipelines/kandinsky/test_kandinsky_inpaint.py | 4 +++- 7 files changed, 36 insertions(+), 18 deletions(-) diff --git a/.github/workflows/pr_modular_tests.yml b/.github/workflows/pr_modular_tests.yml index e01345e325..c6e87e642d 100644 --- a/.github/workflows/pr_modular_tests.yml +++ b/.github/workflows/pr_modular_tests.yml @@ -110,8 +110,9 @@ jobs: run: | python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" python -m uv pip install -e [quality,test] - pip uninstall transformers -y && python -m uv pip install -U transformers@git+https://github.com/huggingface/transformers.git --no-deps - pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git --no-deps + # Stopping this update temporarily until the Hub RC is fully shipped and integrated. + # pip uninstall transformers -y && python -m uv pip install -U transformers@git+https://github.com/huggingface/transformers.git --no-deps + # pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git --no-deps - name: Environment run: | diff --git a/.github/workflows/pr_tests.yml b/.github/workflows/pr_tests.yml index 34a344528e..ebfe9f442f 100644 --- a/.github/workflows/pr_tests.yml +++ b/.github/workflows/pr_tests.yml @@ -116,8 +116,9 @@ jobs: run: | python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" python -m uv pip install -e [quality,test] - pip uninstall transformers -y && python -m uv pip install -U transformers@git+https://github.com/huggingface/transformers.git --no-deps - pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git --no-deps + # Stopping this update temporarily until the Hub RC is fully shipped and integrated. + # pip uninstall transformers -y && python -m uv pip install -U transformers@git+https://github.com/huggingface/transformers.git --no-deps + # pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git --no-deps - name: Environment run: | @@ -253,9 +254,10 @@ jobs: python -m uv pip install -e [quality,test] # TODO (sayakpaul, DN6): revisit `--no-deps` python -m pip install -U peft@git+https://github.com/huggingface/peft.git --no-deps - python -m uv pip install -U transformers@git+https://github.com/huggingface/transformers.git --no-deps - python -m uv pip install -U tokenizers - pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git --no-deps + # Stopping this update temporarily until the Hub RC is fully shipped and integrated. + # python -m uv pip install -U transformers@git+https://github.com/huggingface/transformers.git --no-deps + # python -m uv pip install -U tokenizers + # pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git --no-deps - name: Environment run: | diff --git a/.github/workflows/pr_tests_gpu.yml b/.github/workflows/pr_tests_gpu.yml index 45294c89fe..1a8d5f6b81 100644 --- a/.github/workflows/pr_tests_gpu.yml +++ b/.github/workflows/pr_tests_gpu.yml @@ -132,8 +132,9 @@ jobs: run: | python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" python -m uv pip install -e [quality,test] - pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git - pip uninstall transformers -y && python -m uv pip install -U transformers@git+https://github.com/huggingface/transformers.git --no-deps + # Stopping this update temporarily until the Hub RC is fully shipped and integrated. + # pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git + # pip uninstall transformers -y && python -m uv pip install -U transformers@git+https://github.com/huggingface/transformers.git --no-deps - name: Environment run: | @@ -203,8 +204,9 @@ jobs: python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" python -m uv pip install -e [quality,test] python -m uv pip install peft@git+https://github.com/huggingface/peft.git - pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git - pip uninstall transformers -y && python -m uv pip install -U transformers@git+https://github.com/huggingface/transformers.git --no-deps + # Stopping this update temporarily until the Hub RC is fully shipped and integrated. + # pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git + # pip uninstall transformers -y && python -m uv pip install -U transformers@git+https://github.com/huggingface/transformers.git --no-deps - name: Environment run: | @@ -266,7 +268,8 @@ jobs: - name: Install dependencies run: | python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - pip uninstall transformers -y && python -m uv pip install -U transformers@git+https://github.com/huggingface/transformers.git --no-deps + # Stopping this update temporarily until the Hub RC is fully shipped and integrated. + # pip uninstall transformers -y && python -m uv pip install -U transformers@git+https://github.com/huggingface/transformers.git --no-deps python -m uv pip install -e [quality,test,training] - name: Environment diff --git a/tests/pipelines/kandinsky/test_kandinsky.py b/tests/pipelines/kandinsky/test_kandinsky.py index 9fa39b1bf5..6207e71df8 100644 --- a/tests/pipelines/kandinsky/test_kandinsky.py +++ b/tests/pipelines/kandinsky/test_kandinsky.py @@ -218,7 +218,9 @@ class KandinskyPipelineFastTests(PipelineTesterMixin, unittest.TestCase): return dummy.get_dummy_inputs(device=device, seed=seed) @pytest.mark.xfail( - condition=is_transformers_version(">=", "4.56.2"), reason="Latest transformers changes the slices", strict=True + condition=is_transformers_version(">=", "4.56.2"), + reason="Latest transformers changes the slices", + strict=False, ) def test_kandinsky(self): device = "cpu" diff --git a/tests/pipelines/kandinsky/test_kandinsky_combined.py b/tests/pipelines/kandinsky/test_kandinsky_combined.py index ca80461d87..eba8976597 100644 --- a/tests/pipelines/kandinsky/test_kandinsky_combined.py +++ b/tests/pipelines/kandinsky/test_kandinsky_combined.py @@ -76,7 +76,9 @@ class KandinskyPipelineCombinedFastTests(PipelineTesterMixin, unittest.TestCase) return inputs @pytest.mark.xfail( - condition=is_transformers_version(">=", "4.56.2"), reason="Latest transformers changes the slices", strict=True + condition=is_transformers_version(">=", "4.56.2"), + reason="Latest transformers changes the slices", + strict=False, ) def test_kandinsky(self): device = "cpu" @@ -187,7 +189,9 @@ class KandinskyPipelineImg2ImgCombinedFastTests(PipelineTesterMixin, unittest.Te return inputs @pytest.mark.xfail( - condition=is_transformers_version(">=", "4.56.2"), reason="Latest transformers changes the slices", strict=True + condition=is_transformers_version(">=", "4.56.2"), + reason="Latest transformers changes the slices", + strict=False, ) def test_kandinsky(self): device = "cpu" @@ -301,7 +305,9 @@ class KandinskyPipelineInpaintCombinedFastTests(PipelineTesterMixin, unittest.Te return inputs @pytest.mark.xfail( - condition=is_transformers_version(">=", "4.56.2"), reason="Latest transformers changes the slices", strict=True + condition=is_transformers_version(">=", "4.56.2"), + reason="Latest transformers changes the slices", + strict=False, ) def test_kandinsky(self): device = "cpu" diff --git a/tests/pipelines/kandinsky/test_kandinsky_img2img.py b/tests/pipelines/kandinsky/test_kandinsky_img2img.py index 6bcd9587f2..6d1b43a24f 100644 --- a/tests/pipelines/kandinsky/test_kandinsky_img2img.py +++ b/tests/pipelines/kandinsky/test_kandinsky_img2img.py @@ -240,7 +240,9 @@ class KandinskyImg2ImgPipelineFastTests(PipelineTesterMixin, unittest.TestCase): return dummies.get_dummy_inputs(device=device, seed=seed) @pytest.mark.xfail( - condition=is_transformers_version(">=", "4.56.2"), reason="Latest transformers changes the slices", strict=True + condition=is_transformers_version(">=", "4.56.2"), + reason="Latest transformers changes the slices", + strict=False, ) def test_kandinsky_img2img(self): device = "cpu" diff --git a/tests/pipelines/kandinsky/test_kandinsky_inpaint.py b/tests/pipelines/kandinsky/test_kandinsky_inpaint.py index 6383ca71ef..e2f4aa2a4f 100644 --- a/tests/pipelines/kandinsky/test_kandinsky_inpaint.py +++ b/tests/pipelines/kandinsky/test_kandinsky_inpaint.py @@ -234,7 +234,9 @@ class KandinskyInpaintPipelineFastTests(PipelineTesterMixin, unittest.TestCase): return dummies.get_dummy_inputs(device=device, seed=seed) @pytest.mark.xfail( - condition=is_transformers_version(">=", "4.56.2"), reason="Latest transformers changes the slices", strict=True + condition=is_transformers_version(">=", "4.56.2"), + reason="Latest transformers changes the slices", + strict=False, ) def test_kandinsky_inpaint(self): device = "cpu" From 9c0944581a386736bc808e68d7dfb52d8cf1790e Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Fri, 26 Sep 2025 21:50:16 +0530 Subject: [PATCH 029/123] [docs] slight edits to the attention backends docs. (#12394) * slight edits to the attention backends docs. * Update docs/source/en/optimization/attention_backends.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> --------- Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> --- docs/source/en/optimization/attention_backends.md | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/docs/source/en/optimization/attention_backends.md b/docs/source/en/optimization/attention_backends.md index 04c8b4ba92..e603878a63 100644 --- a/docs/source/en/optimization/attention_backends.md +++ b/docs/source/en/optimization/attention_backends.md @@ -11,7 +11,7 @@ specific language governing permissions and limitations under the License. --> # Attention backends -> [!TIP] +> [!NOTE] > The attention dispatcher is an experimental feature. Please open an issue if you have any feedback or encounter any problems. Diffusers provides several optimized attention algorithms that are more memory and computationally efficient through it's *attention dispatcher*. The dispatcher acts as a router for managing and switching between different attention implementations and provides a unified interface for interacting with them. @@ -33,7 +33,7 @@ The [`~ModelMixin.set_attention_backend`] method iterates through all the module The example below demonstrates how to enable the `_flash_3_hub` implementation for FlashAttention-3 from the [kernel](https://github.com/huggingface/kernels) library, which allows you to instantly use optimized compute kernels from the Hub without requiring any setup. -> [!TIP] +> [!NOTE] > FlashAttention-3 is not supported for non-Hopper architectures, in which case, use FlashAttention with `set_attention_backend("flash")`. ```py @@ -78,10 +78,16 @@ with attention_backend("_flash_3_hub"): image = pipeline(prompt).images[0] ``` +> [!TIP] +> Most attention backends support `torch.compile` without graph breaks and can be used to further speed up inference. + ## Available backends Refer to the table below for a complete list of available attention backends and their variants. +
+Expand + | Backend Name | Family | Description | |--------------|--------|-------------| | `native` | [PyTorch native](https://docs.pytorch.org/docs/stable/generated/torch.nn.attention.SDPBackend.html#torch.nn.attention.SDPBackend) | Default backend using PyTorch's scaled_dot_product_attention | @@ -104,3 +110,5 @@ Refer to the table below for a complete list of available attention backends and | `_sage_qk_int8_pv_fp16_cuda` | [SageAttention](https://github.com/thu-ml/SageAttention) | INT8 QK + FP16 PV (CUDA) | | `_sage_qk_int8_pv_fp16_triton` | [SageAttention](https://github.com/thu-ml/SageAttention) | INT8 QK + FP16 PV (Triton) | | `xformers` | [xFormers](https://github.com/facebookresearch/xformers) | Memory-efficient attention | + +
\ No newline at end of file From 041501aea92919c9c7f36e189fc9cf7d865ebb96 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Fri, 26 Sep 2025 22:38:43 +0530 Subject: [PATCH 030/123] [docs] remove docstrings from repeated methods in `lora_pipeline.py` (#12393) * start unbloating docstrings (save_lora_weights). * load_lora_weights() * lora_state_dict * fuse_lora * unfuse_lora * load_lora_into_transformer --- src/diffusers/loaders/lora_pipeline.py | 2192 +----------------------- 1 file changed, 86 insertions(+), 2106 deletions(-) diff --git a/src/diffusers/loaders/lora_pipeline.py b/src/diffusers/loaders/lora_pipeline.py index 8060b519f1..65bdae6920 100644 --- a/src/diffusers/loaders/lora_pipeline.py +++ b/src/diffusers/loaders/lora_pipeline.py @@ -621,33 +621,7 @@ class StableDiffusionXLLoraLoaderMixin(LoraBaseMixin): **kwargs, ): """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.unet` and - `self.text_encoder`. - - All kwargs are forwarded to `self.lora_state_dict`. - - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is - loaded. - - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details on how the state dict is - loaded into `self.unet`. - - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_text_encoder`] for more details on how the state - dict is loaded into `self.text_encoder`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for more details. """ if not USE_PEFT_BACKEND: raise ValueError("PEFT backend is required for this method.") @@ -967,35 +941,7 @@ class StableDiffusionXLLoraLoaderMixin(LoraBaseMixin): text_encoder_2_lora_adapter_metadata=None, ): r""" - Save the LoRA parameters corresponding to the UNet and text encoder. - - Arguments: - save_directory (`str` or `os.PathLike`): - Directory to save LoRA parameters to. Will be created if it doesn't exist. - unet_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): - State dict of the LoRA layers corresponding to the `unet`. - text_encoder_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): - State dict of the LoRA layers corresponding to the `text_encoder`. Must explicitly pass the text - encoder LoRA state dict because it comes from 🤗 Transformers. - text_encoder_2_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): - State dict of the LoRA layers corresponding to the `text_encoder_2`. Must explicitly pass the text - encoder LoRA state dict because it comes from 🤗 Transformers. - is_main_process (`bool`, *optional*, defaults to `True`): - Whether the process calling this is the main process or not. Useful during distributed training and you - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - unet_lora_adapter_metadata: - LoRA adapter metadata associated with the unet to be serialized with the state dict. - text_encoder_lora_adapter_metadata: - LoRA adapter metadata associated with the text encoder to be serialized with the state dict. - text_encoder_2_lora_adapter_metadata: - LoRA adapter metadata associated with the second text encoder to be serialized with the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. """ lora_layers = {} lora_metadata = {} @@ -1036,35 +982,7 @@ class StableDiffusionXLLoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into. - lora_scale (`float`, defaults to 1.0): - Controls how much to influence the outputs with the LoRA parameters. - safe_fusing (`bool`, defaults to `False`): - Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them. - adapter_names (`List[str]`, *optional*): - Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` + See [`~loaders.StableDiffusionLoraLoaderMixin.fuse_lora`] for more details. """ super().fuse_lora( components=components, @@ -1076,21 +994,7 @@ class StableDiffusionXLLoraLoaderMixin(LoraBaseMixin): def unfuse_lora(self, components: List[str] = ["unet", "text_encoder", "text_encoder_2"], **kwargs): r""" - Reverses the effect of - [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - - This is an experimental API. - - - - Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. - unfuse_unet (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters. - unfuse_text_encoder (`bool`, defaults to `True`): - Whether to unfuse the text encoder LoRA parameters. If the text encoder wasn't monkey-patched with the - LoRA parameters then it won't have any effect. + See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ super().unfuse_lora(components=components, **kwargs) @@ -1116,51 +1020,7 @@ class SD3LoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Return state dict for lora weights and the network alphas. - - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - A [torch state - dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). - - cache_dir (`Union[str, os.PathLike]`, *optional*): - Path to a directory where a downloaded pretrained model configuration is cached if the standard cache - is not used. - force_download (`bool`, *optional*, defaults to `False`): - Whether or not to force the (re-)download of the model weights and configuration files, overriding the - cached versions if they exist. - - proxies (`Dict[str, str]`, *optional*): - A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', - 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - subfolder (`str`, *optional*, defaults to `""`): - The subfolder location of a model file within a larger model repository on the Hub or locally. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. - + See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details. """ # Load the main state dict first which has the LoRA layers for either of # transformer and text encoder or both. @@ -1214,30 +1074,7 @@ class SD3LoraLoaderMixin(LoraBaseMixin): **kwargs, ): """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.unet` and - `self.text_encoder`. - - All kwargs are forwarded to `self.lora_state_dict`. - - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is - loaded. - - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state - dict is loaded into `self.transformer`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for more details. """ if not USE_PEFT_BACKEND: raise ValueError("PEFT backend is required for this method.") @@ -1306,26 +1143,7 @@ class SD3LoraLoaderMixin(LoraBaseMixin): metadata=None, ): """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - A standard state dict containing the lora layer parameters. The keys can either be indexed directly - into the unet or prefixed with an additional `unet` which can be used to distinguish between text - encoder lora layers. - transformer (`SD3Transformer2DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details. """ if low_cpu_mem_usage and is_peft_version("<", "0.13.0"): raise ValueError( @@ -1420,35 +1238,7 @@ class SD3LoraLoaderMixin(LoraBaseMixin): text_encoder_2_lora_adapter_metadata=None, ): r""" - Save the LoRA parameters corresponding to the UNet and text encoder. - - Arguments: - save_directory (`str` or `os.PathLike`): - Directory to save LoRA parameters to. Will be created if it doesn't exist. - transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): - State dict of the LoRA layers corresponding to the `transformer`. - text_encoder_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): - State dict of the LoRA layers corresponding to the `text_encoder`. Must explicitly pass the text - encoder LoRA state dict because it comes from 🤗 Transformers. - text_encoder_2_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): - State dict of the LoRA layers corresponding to the `text_encoder_2`. Must explicitly pass the text - encoder LoRA state dict because it comes from 🤗 Transformers. - is_main_process (`bool`, *optional*, defaults to `True`): - Whether the process calling this is the main process or not. Useful during distributed training and you - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - transformer_lora_adapter_metadata: - LoRA adapter metadata associated with the transformer to be serialized with the state dict. - text_encoder_lora_adapter_metadata: - LoRA adapter metadata associated with the text encoder to be serialized with the state dict. - text_encoder_2_lora_adapter_metadata: - LoRA adapter metadata associated with the second text encoder to be serialized with the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. """ lora_layers = {} lora_metadata = {} @@ -1490,35 +1280,7 @@ class SD3LoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into. - lora_scale (`float`, defaults to 1.0): - Controls how much to influence the outputs with the LoRA parameters. - safe_fusing (`bool`, defaults to `False`): - Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them. - adapter_names (`List[str]`, *optional*): - Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` + See [`~loaders.StableDiffusionLoraLoaderMixin.fuse_lora`] for more details. """ super().fuse_lora( components=components, @@ -1531,21 +1293,7 @@ class SD3LoraLoaderMixin(LoraBaseMixin): # Copied from diffusers.loaders.lora_pipeline.StableDiffusionXLLoraLoaderMixin.unfuse_lora with unet->transformer def unfuse_lora(self, components: List[str] = ["transformer", "text_encoder", "text_encoder_2"], **kwargs): r""" - Reverses the effect of - [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - - This is an experimental API. - - - - Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. - unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters. - unfuse_text_encoder (`bool`, defaults to `True`): - Whether to unfuse the text encoder LoRA parameters. If the text encoder wasn't monkey-patched with the - LoRA parameters then it won't have any effect. + See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ super().unfuse_lora(components=components, **kwargs) @@ -1567,51 +1315,7 @@ class AuraFlowLoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Return state dict for lora weights and the network alphas. - - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - A [torch state - dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). - - cache_dir (`Union[str, os.PathLike]`, *optional*): - Path to a directory where a downloaded pretrained model configuration is cached if the standard cache - is not used. - force_download (`bool`, *optional*, defaults to `False`): - Whether or not to force the (re-)download of the model weights and configuration files, overriding the - cached versions if they exist. - - proxies (`Dict[str, str]`, *optional*): - A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', - 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - subfolder (`str`, *optional*, defaults to `""`): - The subfolder location of a model file within a larger model repository on the Hub or locally. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. - + See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details. """ # Load the main state dict first which has the LoRA layers for either of # transformer and text encoder or both. @@ -1666,25 +1370,7 @@ class AuraFlowLoraLoaderMixin(LoraBaseMixin): **kwargs, ): """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and - `self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See - [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded. - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state - dict is loaded into `self.transformer`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for more details. """ if not USE_PEFT_BACKEND: raise ValueError("PEFT backend is required for this method.") @@ -1730,26 +1416,7 @@ class AuraFlowLoraLoaderMixin(LoraBaseMixin): metadata=None, ): """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - A standard state dict containing the lora layer parameters. The keys can either be indexed directly - into the unet or prefixed with an additional `unet` which can be used to distinguish between text - encoder lora layers. - transformer (`AuraFlowTransformer2DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details. """ if low_cpu_mem_usage and is_peft_version("<", "0.13.0"): raise ValueError( @@ -1781,25 +1448,7 @@ class AuraFlowLoraLoaderMixin(LoraBaseMixin): transformer_lora_adapter_metadata: Optional[dict] = None, ): r""" - Save the LoRA parameters corresponding to the transformer. - - Arguments: - save_directory (`str` or `os.PathLike`): - Directory to save LoRA parameters to. Will be created if it doesn't exist. - transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): - State dict of the LoRA layers corresponding to the `transformer`. - is_main_process (`bool`, *optional*, defaults to `True`): - Whether the process calling this is the main process or not. Useful during distributed training and you - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - transformer_lora_adapter_metadata: - LoRA adapter metadata associated with the transformer to be serialized with the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. """ lora_layers = {} lora_metadata = {} @@ -1831,35 +1480,7 @@ class AuraFlowLoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into. - lora_scale (`float`, defaults to 1.0): - Controls how much to influence the outputs with the LoRA parameters. - safe_fusing (`bool`, defaults to `False`): - Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them. - adapter_names (`List[str]`, *optional*): - Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` + See [`~loaders.StableDiffusionLoraLoaderMixin.fuse_lora`] for more details. """ super().fuse_lora( components=components, @@ -1872,18 +1493,7 @@ class AuraFlowLoraLoaderMixin(LoraBaseMixin): # Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.unfuse_lora def unfuse_lora(self, components: List[str] = ["transformer", "text_encoder"], **kwargs): r""" - Reverses the effect of - [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - - This is an experimental API. - - - - Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. - unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters. + See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ super().unfuse_lora(components=components, **kwargs) @@ -1910,50 +1520,7 @@ class FluxLoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Return state dict for lora weights and the network alphas. - - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - A [torch state - dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). - - cache_dir (`Union[str, os.PathLike]`, *optional*): - Path to a directory where a downloaded pretrained model configuration is cached if the standard cache - is not used. - force_download (`bool`, *optional*, defaults to `False`): - Whether or not to force the (re-)download of the model weights and configuration files, overriding the - cached versions if they exist. - - proxies (`Dict[str, str]`, *optional*): - A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', - 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - subfolder (`str`, *optional*, defaults to `""`): - The subfolder location of a model file within a larger model repository on the Hub or locally. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details. """ # Load the main state dict first which has the LoRA layers for either of # transformer and text encoder or both. @@ -2207,30 +1774,7 @@ class FluxLoraLoaderMixin(LoraBaseMixin): hotswap: bool = False, ): """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - A standard state dict containing the lora layer parameters. The keys can either be indexed directly - into the unet or prefixed with an additional `unet` which can be used to distinguish between text - encoder lora layers. - network_alphas (`Dict[str, float]`): - The value of the network alpha used for stable learning and preventing underflow. This value has the - same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this - link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning). - transformer (`FluxTransformer2DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details. """ if low_cpu_mem_usage and not is_peft_version(">=", "0.13.1"): raise ValueError( @@ -2435,35 +1979,7 @@ class FluxLoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into. - lora_scale (`float`, defaults to 1.0): - Controls how much to influence the outputs with the LoRA parameters. - safe_fusing (`bool`, defaults to `False`): - Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them. - adapter_names (`List[str]`, *optional*): - Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` + See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details. """ transformer = getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer @@ -2806,30 +2322,7 @@ class AmusedLoraLoaderMixin(StableDiffusionLoraLoaderMixin): hotswap: bool = False, ): """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - A standard state dict containing the lora layer parameters. The keys can either be indexed directly - into the unet or prefixed with an additional `unet` which can be used to distinguish between text - encoder lora layers. - network_alphas (`Dict[str, float]`): - The value of the network alpha used for stable learning and preventing underflow. This value has the - same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this - link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning). - transformer (`UVit2DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details. """ if low_cpu_mem_usage and not is_peft_version(">=", "0.13.1"): raise ValueError( @@ -2979,51 +2472,7 @@ class CogVideoXLoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Return state dict for lora weights and the network alphas. - - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - A [torch state - dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). - - cache_dir (`Union[str, os.PathLike]`, *optional*): - Path to a directory where a downloaded pretrained model configuration is cached if the standard cache - is not used. - force_download (`bool`, *optional*, defaults to `False`): - Whether or not to force the (re-)download of the model weights and configuration files, overriding the - cached versions if they exist. - - proxies (`Dict[str, str]`, *optional*): - A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', - 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - subfolder (`str`, *optional*, defaults to `""`): - The subfolder location of a model file within a larger model repository on the Hub or locally. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. - + See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details. """ # Load the main state dict first which has the LoRA layers for either of # transformer and text encoder or both. @@ -3077,25 +2526,7 @@ class CogVideoXLoraLoaderMixin(LoraBaseMixin): **kwargs, ): """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and - `self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See - [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded. - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state - dict is loaded into `self.transformer`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for more details. """ if not USE_PEFT_BACKEND: raise ValueError("PEFT backend is required for this method.") @@ -3141,26 +2572,7 @@ class CogVideoXLoraLoaderMixin(LoraBaseMixin): metadata=None, ): """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - A standard state dict containing the lora layer parameters. The keys can either be indexed directly - into the unet or prefixed with an additional `unet` which can be used to distinguish between text - encoder lora layers. - transformer (`CogVideoXTransformer3DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details. """ if low_cpu_mem_usage and is_peft_version("<", "0.13.0"): raise ValueError( @@ -3180,7 +2592,6 @@ class CogVideoXLoraLoaderMixin(LoraBaseMixin): ) @classmethod - # Adapted from diffusers.loaders.lora_pipeline.StableDiffusionLoraLoaderMixin.save_lora_weights without support for text encoder def save_lora_weights( cls, save_directory: Union[str, os.PathLike], @@ -3192,25 +2603,7 @@ class CogVideoXLoraLoaderMixin(LoraBaseMixin): transformer_lora_adapter_metadata: Optional[dict] = None, ): r""" - Save the LoRA parameters corresponding to the transformer. - - Arguments: - save_directory (`str` or `os.PathLike`): - Directory to save LoRA parameters to. Will be created if it doesn't exist. - transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): - State dict of the LoRA layers corresponding to the `transformer`. - is_main_process (`bool`, *optional*, defaults to `True`): - Whether the process calling this is the main process or not. Useful during distributed training and you - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - transformer_lora_adapter_metadata: - LoRA adapter metadata associated with the transformer to be serialized with the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. """ lora_layers = {} lora_metadata = {} @@ -3241,35 +2634,7 @@ class CogVideoXLoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into. - lora_scale (`float`, defaults to 1.0): - Controls how much to influence the outputs with the LoRA parameters. - safe_fusing (`bool`, defaults to `False`): - Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them. - adapter_names (`List[str]`, *optional*): - Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` + See [`~loaders.StableDiffusionLoraLoaderMixin.fuse_lora`] for more details. """ super().fuse_lora( components=components, @@ -3281,18 +2646,7 @@ class CogVideoXLoraLoaderMixin(LoraBaseMixin): def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): r""" - Reverses the effect of - [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - - This is an experimental API. - - - - Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. - unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters. + See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ super().unfuse_lora(components=components, **kwargs) @@ -3314,51 +2668,7 @@ class Mochi1LoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Return state dict for lora weights and the network alphas. - - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - A [torch state - dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). - - cache_dir (`Union[str, os.PathLike]`, *optional*): - Path to a directory where a downloaded pretrained model configuration is cached if the standard cache - is not used. - force_download (`bool`, *optional*, defaults to `False`): - Whether or not to force the (re-)download of the model weights and configuration files, overriding the - cached versions if they exist. - - proxies (`Dict[str, str]`, *optional*): - A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', - 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - subfolder (`str`, *optional*, defaults to `""`): - The subfolder location of a model file within a larger model repository on the Hub or locally. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. - + See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details. """ # Load the main state dict first which has the LoRA layers for either of # transformer and text encoder or both. @@ -3413,25 +2723,7 @@ class Mochi1LoraLoaderMixin(LoraBaseMixin): **kwargs, ): """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and - `self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See - [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded. - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state - dict is loaded into `self.transformer`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for more details. """ if not USE_PEFT_BACKEND: raise ValueError("PEFT backend is required for this method.") @@ -3477,26 +2769,7 @@ class Mochi1LoraLoaderMixin(LoraBaseMixin): metadata=None, ): """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - A standard state dict containing the lora layer parameters. The keys can either be indexed directly - into the unet or prefixed with an additional `unet` which can be used to distinguish between text - encoder lora layers. - transformer (`MochiTransformer3DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details. """ if low_cpu_mem_usage and is_peft_version("<", "0.13.0"): raise ValueError( @@ -3528,25 +2801,7 @@ class Mochi1LoraLoaderMixin(LoraBaseMixin): transformer_lora_adapter_metadata: Optional[dict] = None, ): r""" - Save the LoRA parameters corresponding to the transformer. - - Arguments: - save_directory (`str` or `os.PathLike`): - Directory to save LoRA parameters to. Will be created if it doesn't exist. - transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): - State dict of the LoRA layers corresponding to the `transformer`. - is_main_process (`bool`, *optional*, defaults to `True`): - Whether the process calling this is the main process or not. Useful during distributed training and you - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - transformer_lora_adapter_metadata: - LoRA adapter metadata associated with the transformer to be serialized with the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. """ lora_layers = {} lora_metadata = {} @@ -3578,35 +2833,7 @@ class Mochi1LoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into. - lora_scale (`float`, defaults to 1.0): - Controls how much to influence the outputs with the LoRA parameters. - safe_fusing (`bool`, defaults to `False`): - Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them. - adapter_names (`List[str]`, *optional*): - Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` + See [`~loaders.StableDiffusionLoraLoaderMixin.fuse_lora`] for more details. """ super().fuse_lora( components=components, @@ -3619,18 +2846,7 @@ class Mochi1LoraLoaderMixin(LoraBaseMixin): # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): r""" - Reverses the effect of - [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - - This is an experimental API. - - - - Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. - unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters. + See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ super().unfuse_lora(components=components, **kwargs) @@ -3651,50 +2867,7 @@ class LTXVideoLoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Return state dict for lora weights and the network alphas. - - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - A [torch state - dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). - - cache_dir (`Union[str, os.PathLike]`, *optional*): - Path to a directory where a downloaded pretrained model configuration is cached if the standard cache - is not used. - force_download (`bool`, *optional*, defaults to `False`): - Whether or not to force the (re-)download of the model weights and configuration files, overriding the - cached versions if they exist. - - proxies (`Dict[str, str]`, *optional*): - A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', - 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - subfolder (`str`, *optional*, defaults to `""`): - The subfolder location of a model file within a larger model repository on the Hub or locally. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details. """ # Load the main state dict first which has the LoRA layers for either of # transformer and text encoder or both. @@ -3753,25 +2926,7 @@ class LTXVideoLoraLoaderMixin(LoraBaseMixin): **kwargs, ): """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and - `self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See - [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded. - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state - dict is loaded into `self.transformer`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for more details. """ if not USE_PEFT_BACKEND: raise ValueError("PEFT backend is required for this method.") @@ -3817,26 +2972,7 @@ class LTXVideoLoraLoaderMixin(LoraBaseMixin): metadata=None, ): """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - A standard state dict containing the lora layer parameters. The keys can either be indexed directly - into the unet or prefixed with an additional `unet` which can be used to distinguish between text - encoder lora layers. - transformer (`LTXVideoTransformer3DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details. """ if low_cpu_mem_usage and is_peft_version("<", "0.13.0"): raise ValueError( @@ -3868,25 +3004,7 @@ class LTXVideoLoraLoaderMixin(LoraBaseMixin): transformer_lora_adapter_metadata: Optional[dict] = None, ): r""" - Save the LoRA parameters corresponding to the transformer. - - Arguments: - save_directory (`str` or `os.PathLike`): - Directory to save LoRA parameters to. Will be created if it doesn't exist. - transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): - State dict of the LoRA layers corresponding to the `transformer`. - is_main_process (`bool`, *optional*, defaults to `True`): - Whether the process calling this is the main process or not. Useful during distributed training and you - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - transformer_lora_adapter_metadata: - LoRA adapter metadata associated with the transformer to be serialized with the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. """ lora_layers = {} lora_metadata = {} @@ -3918,35 +3036,7 @@ class LTXVideoLoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into. - lora_scale (`float`, defaults to 1.0): - Controls how much to influence the outputs with the LoRA parameters. - safe_fusing (`bool`, defaults to `False`): - Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them. - adapter_names (`List[str]`, *optional*): - Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` + See [`~loaders.StableDiffusionLoraLoaderMixin.fuse_lora`] for more details. """ super().fuse_lora( components=components, @@ -3959,18 +3049,7 @@ class LTXVideoLoraLoaderMixin(LoraBaseMixin): # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): r""" - Reverses the effect of - [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - - This is an experimental API. - - - - Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. - unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters. + See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ super().unfuse_lora(components=components, **kwargs) @@ -3992,51 +3071,7 @@ class SanaLoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Return state dict for lora weights and the network alphas. - - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - A [torch state - dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). - - cache_dir (`Union[str, os.PathLike]`, *optional*): - Path to a directory where a downloaded pretrained model configuration is cached if the standard cache - is not used. - force_download (`bool`, *optional*, defaults to `False`): - Whether or not to force the (re-)download of the model weights and configuration files, overriding the - cached versions if they exist. - - proxies (`Dict[str, str]`, *optional*): - A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', - 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - subfolder (`str`, *optional*, defaults to `""`): - The subfolder location of a model file within a larger model repository on the Hub or locally. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. - + See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details. """ # Load the main state dict first which has the LoRA layers for either of # transformer and text encoder or both. @@ -4091,25 +3126,7 @@ class SanaLoraLoaderMixin(LoraBaseMixin): **kwargs, ): """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and - `self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See - [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded. - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state - dict is loaded into `self.transformer`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for more details. """ if not USE_PEFT_BACKEND: raise ValueError("PEFT backend is required for this method.") @@ -4155,26 +3172,7 @@ class SanaLoraLoaderMixin(LoraBaseMixin): metadata=None, ): """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - A standard state dict containing the lora layer parameters. The keys can either be indexed directly - into the unet or prefixed with an additional `unet` which can be used to distinguish between text - encoder lora layers. - transformer (`SanaTransformer2DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details. """ if low_cpu_mem_usage and is_peft_version("<", "0.13.0"): raise ValueError( @@ -4206,25 +3204,7 @@ class SanaLoraLoaderMixin(LoraBaseMixin): transformer_lora_adapter_metadata: Optional[dict] = None, ): r""" - Save the LoRA parameters corresponding to the transformer. - - Arguments: - save_directory (`str` or `os.PathLike`): - Directory to save LoRA parameters to. Will be created if it doesn't exist. - transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): - State dict of the LoRA layers corresponding to the `transformer`. - is_main_process (`bool`, *optional*, defaults to `True`): - Whether the process calling this is the main process or not. Useful during distributed training and you - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - transformer_lora_adapter_metadata: - LoRA adapter metadata associated with the transformer to be serialized with the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. """ lora_layers = {} lora_metadata = {} @@ -4256,35 +3236,7 @@ class SanaLoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into. - lora_scale (`float`, defaults to 1.0): - Controls how much to influence the outputs with the LoRA parameters. - safe_fusing (`bool`, defaults to `False`): - Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them. - adapter_names (`List[str]`, *optional*): - Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` + See [`~loaders.StableDiffusionLoraLoaderMixin.fuse_lora`] for more details. """ super().fuse_lora( components=components, @@ -4297,18 +3249,7 @@ class SanaLoraLoaderMixin(LoraBaseMixin): # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): r""" - Reverses the effect of - [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - - This is an experimental API. - - - - Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. - unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters. + See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ super().unfuse_lora(components=components, **kwargs) @@ -4329,50 +3270,7 @@ class HunyuanVideoLoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Return state dict for lora weights and the network alphas. - - - - We support loading original format HunyuanVideo LoRA checkpoints. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - A [torch state - dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). - - cache_dir (`Union[str, os.PathLike]`, *optional*): - Path to a directory where a downloaded pretrained model configuration is cached if the standard cache - is not used. - force_download (`bool`, *optional*, defaults to `False`): - Whether or not to force the (re-)download of the model weights and configuration files, overriding the - cached versions if they exist. - - proxies (`Dict[str, str]`, *optional*): - A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', - 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - subfolder (`str`, *optional*, defaults to `""`): - The subfolder location of a model file within a larger model repository on the Hub or locally. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details. """ # Load the main state dict first which has the LoRA layers for either of # transformer and text encoder or both. @@ -4431,25 +3329,7 @@ class HunyuanVideoLoraLoaderMixin(LoraBaseMixin): **kwargs, ): """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and - `self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See - [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded. - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state - dict is loaded into `self.transformer`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for more details. """ if not USE_PEFT_BACKEND: raise ValueError("PEFT backend is required for this method.") @@ -4495,26 +3375,7 @@ class HunyuanVideoLoraLoaderMixin(LoraBaseMixin): metadata=None, ): """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - A standard state dict containing the lora layer parameters. The keys can either be indexed directly - into the unet or prefixed with an additional `unet` which can be used to distinguish between text - encoder lora layers. - transformer (`HunyuanVideoTransformer3DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details. """ if low_cpu_mem_usage and is_peft_version("<", "0.13.0"): raise ValueError( @@ -4546,25 +3407,7 @@ class HunyuanVideoLoraLoaderMixin(LoraBaseMixin): transformer_lora_adapter_metadata: Optional[dict] = None, ): r""" - Save the LoRA parameters corresponding to the transformer. - - Arguments: - save_directory (`str` or `os.PathLike`): - Directory to save LoRA parameters to. Will be created if it doesn't exist. - transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): - State dict of the LoRA layers corresponding to the `transformer`. - is_main_process (`bool`, *optional*, defaults to `True`): - Whether the process calling this is the main process or not. Useful during distributed training and you - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - transformer_lora_adapter_metadata: - LoRA adapter metadata associated with the transformer to be serialized with the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. """ lora_layers = {} lora_metadata = {} @@ -4596,35 +3439,7 @@ class HunyuanVideoLoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into. - lora_scale (`float`, defaults to 1.0): - Controls how much to influence the outputs with the LoRA parameters. - safe_fusing (`bool`, defaults to `False`): - Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them. - adapter_names (`List[str]`, *optional*): - Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` + See [`~loaders.StableDiffusionLoraLoaderMixin.fuse_lora`] for more details. """ super().fuse_lora( components=components, @@ -4637,18 +3452,7 @@ class HunyuanVideoLoraLoaderMixin(LoraBaseMixin): # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): r""" - Reverses the effect of - [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - - This is an experimental API. - - - - Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. - unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters. + See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ super().unfuse_lora(components=components, **kwargs) @@ -4669,50 +3473,7 @@ class Lumina2LoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Return state dict for lora weights and the network alphas. - - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - A [torch state - dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). - - cache_dir (`Union[str, os.PathLike]`, *optional*): - Path to a directory where a downloaded pretrained model configuration is cached if the standard cache - is not used. - force_download (`bool`, *optional*, defaults to `False`): - Whether or not to force the (re-)download of the model weights and configuration files, overriding the - cached versions if they exist. - - proxies (`Dict[str, str]`, *optional*): - A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', - 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - subfolder (`str`, *optional*, defaults to `""`): - The subfolder location of a model file within a larger model repository on the Hub or locally. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details. """ # Load the main state dict first which has the LoRA layers for either of # transformer and text encoder or both. @@ -4772,25 +3533,7 @@ class Lumina2LoraLoaderMixin(LoraBaseMixin): **kwargs, ): """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and - `self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See - [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded. - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state - dict is loaded into `self.transformer`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for more details. """ if not USE_PEFT_BACKEND: raise ValueError("PEFT backend is required for this method.") @@ -4836,26 +3579,7 @@ class Lumina2LoraLoaderMixin(LoraBaseMixin): metadata=None, ): """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - A standard state dict containing the lora layer parameters. The keys can either be indexed directly - into the unet or prefixed with an additional `unet` which can be used to distinguish between text - encoder lora layers. - transformer (`Lumina2Transformer2DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details. """ if low_cpu_mem_usage and is_peft_version("<", "0.13.0"): raise ValueError( @@ -4887,25 +3611,7 @@ class Lumina2LoraLoaderMixin(LoraBaseMixin): transformer_lora_adapter_metadata: Optional[dict] = None, ): r""" - Save the LoRA parameters corresponding to the transformer. - - Arguments: - save_directory (`str` or `os.PathLike`): - Directory to save LoRA parameters to. Will be created if it doesn't exist. - transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): - State dict of the LoRA layers corresponding to the `transformer`. - is_main_process (`bool`, *optional*, defaults to `True`): - Whether the process calling this is the main process or not. Useful during distributed training and you - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - transformer_lora_adapter_metadata: - LoRA adapter metadata associated with the transformer to be serialized with the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. """ lora_layers = {} lora_metadata = {} @@ -4937,35 +3643,7 @@ class Lumina2LoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into. - lora_scale (`float`, defaults to 1.0): - Controls how much to influence the outputs with the LoRA parameters. - safe_fusing (`bool`, defaults to `False`): - Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them. - adapter_names (`List[str]`, *optional*): - Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` + See [`~loaders.StableDiffusionLoraLoaderMixin.fuse_lora`] for more details. """ super().fuse_lora( components=components, @@ -4978,18 +3656,7 @@ class Lumina2LoraLoaderMixin(LoraBaseMixin): # Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.unfuse_lora def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): r""" - Reverses the effect of - [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - - This is an experimental API. - - - - Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. - unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters. + See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ super().unfuse_lora(components=components, **kwargs) @@ -5010,50 +3677,7 @@ class WanLoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Return state dict for lora weights and the network alphas. - - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - A [torch state - dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). - - cache_dir (`Union[str, os.PathLike]`, *optional*): - Path to a directory where a downloaded pretrained model configuration is cached if the standard cache - is not used. - force_download (`bool`, *optional*, defaults to `False`): - Whether or not to force the (re-)download of the model weights and configuration files, overriding the - cached versions if they exist. - - proxies (`Dict[str, str]`, *optional*): - A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', - 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - subfolder (`str`, *optional*, defaults to `""`): - The subfolder location of a model file within a larger model repository on the Hub or locally. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details. """ # Load the main state dict first which has the LoRA layers for either of # transformer and text encoder or both. @@ -5159,25 +3783,7 @@ class WanLoraLoaderMixin(LoraBaseMixin): **kwargs, ): """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and - `self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See - [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded. - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state - dict is loaded into `self.transformer`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for more details. """ if not USE_PEFT_BACKEND: raise ValueError("PEFT backend is required for this method.") @@ -5247,26 +3853,7 @@ class WanLoraLoaderMixin(LoraBaseMixin): metadata=None, ): """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - A standard state dict containing the lora layer parameters. The keys can either be indexed directly - into the unet or prefixed with an additional `unet` which can be used to distinguish between text - encoder lora layers. - transformer (`WanTransformer3DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details. """ if low_cpu_mem_usage and is_peft_version("<", "0.13.0"): raise ValueError( @@ -5298,25 +3885,7 @@ class WanLoraLoaderMixin(LoraBaseMixin): transformer_lora_adapter_metadata: Optional[dict] = None, ): r""" - Save the LoRA parameters corresponding to the transformer. - - Arguments: - save_directory (`str` or `os.PathLike`): - Directory to save LoRA parameters to. Will be created if it doesn't exist. - transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): - State dict of the LoRA layers corresponding to the `transformer`. - is_main_process (`bool`, *optional*, defaults to `True`): - Whether the process calling this is the main process or not. Useful during distributed training and you - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - transformer_lora_adapter_metadata: - LoRA adapter metadata associated with the transformer to be serialized with the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. """ lora_layers = {} lora_metadata = {} @@ -5348,35 +3917,7 @@ class WanLoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into. - lora_scale (`float`, defaults to 1.0): - Controls how much to influence the outputs with the LoRA parameters. - safe_fusing (`bool`, defaults to `False`): - Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them. - adapter_names (`List[str]`, *optional*): - Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` + See [`~loaders.StableDiffusionLoraLoaderMixin.fuse_lora`] for more details. """ super().fuse_lora( components=components, @@ -5389,18 +3930,7 @@ class WanLoraLoaderMixin(LoraBaseMixin): # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): r""" - Reverses the effect of - [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - - This is an experimental API. - - - - Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. - unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters. + See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ super().unfuse_lora(components=components, **kwargs) @@ -5422,50 +3952,7 @@ class SkyReelsV2LoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Return state dict for lora weights and the network alphas. - - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - A [torch state - dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). - - cache_dir (`Union[str, os.PathLike]`, *optional*): - Path to a directory where a downloaded pretrained model configuration is cached if the standard cache - is not used. - force_download (`bool`, *optional*, defaults to `False`): - Whether or not to force the (re-)download of the model weights and configuration files, overriding the - cached versions if they exist. - - proxies (`Dict[str, str]`, *optional*): - A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', - 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - subfolder (`str`, *optional*, defaults to `""`): - The subfolder location of a model file within a larger model repository on the Hub or locally. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details. """ # Load the main state dict first which has the LoRA layers for either of # transformer and text encoder or both. @@ -5573,25 +4060,7 @@ class SkyReelsV2LoraLoaderMixin(LoraBaseMixin): **kwargs, ): """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and - `self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See - [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded. - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state - dict is loaded into `self.transformer`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for more details. """ if not USE_PEFT_BACKEND: raise ValueError("PEFT backend is required for this method.") @@ -5661,26 +4130,7 @@ class SkyReelsV2LoraLoaderMixin(LoraBaseMixin): metadata=None, ): """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - A standard state dict containing the lora layer parameters. The keys can either be indexed directly - into the unet or prefixed with an additional `unet` which can be used to distinguish between text - encoder lora layers. - transformer (`SkyReelsV2Transformer3DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details. """ if low_cpu_mem_usage and is_peft_version("<", "0.13.0"): raise ValueError( @@ -5712,25 +4162,7 @@ class SkyReelsV2LoraLoaderMixin(LoraBaseMixin): transformer_lora_adapter_metadata: Optional[dict] = None, ): r""" - Save the LoRA parameters corresponding to the transformer. - - Arguments: - save_directory (`str` or `os.PathLike`): - Directory to save LoRA parameters to. Will be created if it doesn't exist. - transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): - State dict of the LoRA layers corresponding to the `transformer`. - is_main_process (`bool`, *optional*, defaults to `True`): - Whether the process calling this is the main process or not. Useful during distributed training and you - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - transformer_lora_adapter_metadata: - LoRA adapter metadata associated with the transformer to be serialized with the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. """ lora_layers = {} lora_metadata = {} @@ -5762,35 +4194,7 @@ class SkyReelsV2LoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into. - lora_scale (`float`, defaults to 1.0): - Controls how much to influence the outputs with the LoRA parameters. - safe_fusing (`bool`, defaults to `False`): - Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them. - adapter_names (`List[str]`, *optional*): - Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` + See [`~loaders.StableDiffusionLoraLoaderMixin.fuse_lora`] for more details. """ super().fuse_lora( components=components, @@ -5803,18 +4207,7 @@ class SkyReelsV2LoraLoaderMixin(LoraBaseMixin): # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): r""" - Reverses the effect of - [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - - This is an experimental API. - - - - Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. - unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters. + See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ super().unfuse_lora(components=components, **kwargs) @@ -5836,51 +4229,7 @@ class CogView4LoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Return state dict for lora weights and the network alphas. - - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - A [torch state - dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). - - cache_dir (`Union[str, os.PathLike]`, *optional*): - Path to a directory where a downloaded pretrained model configuration is cached if the standard cache - is not used. - force_download (`bool`, *optional*, defaults to `False`): - Whether or not to force the (re-)download of the model weights and configuration files, overriding the - cached versions if they exist. - - proxies (`Dict[str, str]`, *optional*): - A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', - 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - subfolder (`str`, *optional*, defaults to `""`): - The subfolder location of a model file within a larger model repository on the Hub or locally. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. - + See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details. """ # Load the main state dict first which has the LoRA layers for either of # transformer and text encoder or both. @@ -5935,25 +4284,7 @@ class CogView4LoraLoaderMixin(LoraBaseMixin): **kwargs, ): """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and - `self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See - [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded. - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state - dict is loaded into `self.transformer`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for more details. """ if not USE_PEFT_BACKEND: raise ValueError("PEFT backend is required for this method.") @@ -5999,26 +4330,7 @@ class CogView4LoraLoaderMixin(LoraBaseMixin): metadata=None, ): """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - A standard state dict containing the lora layer parameters. The keys can either be indexed directly - into the unet or prefixed with an additional `unet` which can be used to distinguish between text - encoder lora layers. - transformer (`CogView4Transformer2DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details. """ if low_cpu_mem_usage and is_peft_version("<", "0.13.0"): raise ValueError( @@ -6050,25 +4362,7 @@ class CogView4LoraLoaderMixin(LoraBaseMixin): transformer_lora_adapter_metadata: Optional[dict] = None, ): r""" - Save the LoRA parameters corresponding to the transformer. - - Arguments: - save_directory (`str` or `os.PathLike`): - Directory to save LoRA parameters to. Will be created if it doesn't exist. - transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): - State dict of the LoRA layers corresponding to the `transformer`. - is_main_process (`bool`, *optional*, defaults to `True`): - Whether the process calling this is the main process or not. Useful during distributed training and you - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - transformer_lora_adapter_metadata: - LoRA adapter metadata associated with the transformer to be serialized with the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. """ lora_layers = {} lora_metadata = {} @@ -6100,35 +4394,7 @@ class CogView4LoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into. - lora_scale (`float`, defaults to 1.0): - Controls how much to influence the outputs with the LoRA parameters. - safe_fusing (`bool`, defaults to `False`): - Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them. - adapter_names (`List[str]`, *optional*): - Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` + See [`~loaders.StableDiffusionLoraLoaderMixin.fuse_lora`] for more details. """ super().fuse_lora( components=components, @@ -6141,18 +4407,7 @@ class CogView4LoraLoaderMixin(LoraBaseMixin): # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): r""" - Reverses the effect of - [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - - This is an experimental API. - - - - Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. - unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters. + See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ super().unfuse_lora(components=components, **kwargs) @@ -6173,50 +4428,7 @@ class HiDreamImageLoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Return state dict for lora weights and the network alphas. - - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - A [torch state - dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). - - cache_dir (`Union[str, os.PathLike]`, *optional*): - Path to a directory where a downloaded pretrained model configuration is cached if the standard cache - is not used. - force_download (`bool`, *optional*, defaults to `False`): - Whether or not to force the (re-)download of the model weights and configuration files, overriding the - cached versions if they exist. - - proxies (`Dict[str, str]`, *optional*): - A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', - 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - subfolder (`str`, *optional*, defaults to `""`): - The subfolder location of a model file within a larger model repository on the Hub or locally. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details. """ # Load the main state dict first which has the LoRA layers for either of # transformer and text encoder or both. @@ -6275,25 +4487,7 @@ class HiDreamImageLoraLoaderMixin(LoraBaseMixin): **kwargs, ): """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and - `self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See - [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded. - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state - dict is loaded into `self.transformer`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for more details. """ if not USE_PEFT_BACKEND: raise ValueError("PEFT backend is required for this method.") @@ -6339,26 +4533,7 @@ class HiDreamImageLoraLoaderMixin(LoraBaseMixin): metadata=None, ): """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - A standard state dict containing the lora layer parameters. The keys can either be indexed directly - into the unet or prefixed with an additional `unet` which can be used to distinguish between text - encoder lora layers. - transformer (`HiDreamImageTransformer2DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details. """ if low_cpu_mem_usage and is_peft_version("<", "0.13.0"): raise ValueError( @@ -6390,25 +4565,7 @@ class HiDreamImageLoraLoaderMixin(LoraBaseMixin): transformer_lora_adapter_metadata: Optional[dict] = None, ): r""" - Save the LoRA parameters corresponding to the transformer. - - Arguments: - save_directory (`str` or `os.PathLike`): - Directory to save LoRA parameters to. Will be created if it doesn't exist. - transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): - State dict of the LoRA layers corresponding to the `transformer`. - is_main_process (`bool`, *optional*, defaults to `True`): - Whether the process calling this is the main process or not. Useful during distributed training and you - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - transformer_lora_adapter_metadata: - LoRA adapter metadata associated with the transformer to be serialized with the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. """ lora_layers = {} lora_metadata = {} @@ -6440,35 +4597,7 @@ class HiDreamImageLoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into. - lora_scale (`float`, defaults to 1.0): - Controls how much to influence the outputs with the LoRA parameters. - safe_fusing (`bool`, defaults to `False`): - Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them. - adapter_names (`List[str]`, *optional*): - Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` + See [`~loaders.StableDiffusionLoraLoaderMixin.fuse_lora`] for more details. """ super().fuse_lora( components=components, @@ -6481,18 +4610,7 @@ class HiDreamImageLoraLoaderMixin(LoraBaseMixin): # Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.unfuse_lora def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): r""" - Reverses the effect of - [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - - This is an experimental API. - - - - Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. - unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters. + See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ super().unfuse_lora(components=components, **kwargs) @@ -6513,51 +4631,7 @@ class QwenImageLoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Return state dict for lora weights and the network alphas. - - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - A [torch state - dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). - - cache_dir (`Union[str, os.PathLike]`, *optional*): - Path to a directory where a downloaded pretrained model configuration is cached if the standard cache - is not used. - force_download (`bool`, *optional*, defaults to `False`): - Whether or not to force the (re-)download of the model weights and configuration files, overriding the - cached versions if they exist. - - proxies (`Dict[str, str]`, *optional*): - A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', - 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - subfolder (`str`, *optional*, defaults to `""`): - The subfolder location of a model file within a larger model repository on the Hub or locally. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. - + See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details. """ # Load the main state dict first which has the LoRA layers for either of # transformer and text encoder or both. @@ -6618,25 +4692,7 @@ class QwenImageLoraLoaderMixin(LoraBaseMixin): **kwargs, ): """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and - `self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See - [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded. - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state - dict is loaded into `self.transformer`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for more details. """ if not USE_PEFT_BACKEND: raise ValueError("PEFT backend is required for this method.") @@ -6682,26 +4738,7 @@ class QwenImageLoraLoaderMixin(LoraBaseMixin): metadata=None, ): """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - A standard state dict containing the lora layer parameters. The keys can either be indexed directly - into the unet or prefixed with an additional `unet` which can be used to distinguish between text - encoder lora layers. - transformer (`QwenImageTransformer2DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details. """ if low_cpu_mem_usage and is_peft_version("<", "0.13.0"): raise ValueError( @@ -6733,25 +4770,7 @@ class QwenImageLoraLoaderMixin(LoraBaseMixin): transformer_lora_adapter_metadata: Optional[dict] = None, ): r""" - Save the LoRA parameters corresponding to the transformer. - - Arguments: - save_directory (`str` or `os.PathLike`): - Directory to save LoRA parameters to. Will be created if it doesn't exist. - transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): - State dict of the LoRA layers corresponding to the `transformer`. - is_main_process (`bool`, *optional*, defaults to `True`): - Whether the process calling this is the main process or not. Useful during distributed training and you - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - transformer_lora_adapter_metadata: - LoRA adapter metadata associated with the transformer to be serialized with the state dict. + See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. """ lora_layers = {} lora_metadata = {} @@ -6783,35 +4802,7 @@ class QwenImageLoraLoaderMixin(LoraBaseMixin): **kwargs, ): r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into. - lora_scale (`float`, defaults to 1.0): - Controls how much to influence the outputs with the LoRA parameters. - safe_fusing (`bool`, defaults to `False`): - Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them. - adapter_names (`List[str]`, *optional*): - Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` + See [`~loaders.StableDiffusionLoraLoaderMixin.fuse_lora`] for more details. """ super().fuse_lora( components=components, @@ -6824,18 +4815,7 @@ class QwenImageLoraLoaderMixin(LoraBaseMixin): # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): r""" - Reverses the effect of - [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - - This is an experimental API. - - - - Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. - unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters. + See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ super().unfuse_lora(components=components, **kwargs) From 19085ac8f4947091b1e6b3ca980153eadc12c653 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Mon, 29 Sep 2025 13:08:05 +0530 Subject: [PATCH 031/123] Don't skip Qwen model tests for group offloading with disk (#12382) u[ --- tests/models/test_modeling_common.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/tests/models/test_modeling_common.py b/tests/models/test_modeling_common.py index 3a008edfe1..90ded6a7ec 100644 --- a/tests/models/test_modeling_common.py +++ b/tests/models/test_modeling_common.py @@ -1793,11 +1793,6 @@ class ModelTesterMixin: if not self.model_class._supports_group_offloading: pytest.skip("Model does not support group offloading.") - if self.model_class.__name__ == "QwenImageTransformer2DModel": - pytest.skip( - "QwenImageTransformer2DModel doesn't support group offloading with disk. Needs to be investigated." - ) - def _has_generator_arg(model): sig = inspect.signature(model.forward) params = sig.parameters From 0a151115bbe493de74a4565e57352a0890e94777 Mon Sep 17 00:00:00 2001 From: Akshay Babbar Date: Mon, 29 Sep 2025 14:20:05 +0530 Subject: [PATCH 032/123] Fix #12116: preserve boolean dtype for attention masks in ChromaPipeline (#12263) * fix: preserve boolean dtype for attention masks in ChromaPipeline - Convert attention masks to bool and prevent dtype corruption - Fix both positive and negative mask handling in _get_t5_prompt_embeds - Remove float conversion in _prepare_attention_mask method Fixes #12116 * test: add ChromaPipeline attention mask dtype tests * test: add slow ChromaPipeline attention mask tests * chore: removed comments * refactor: removing redundant type conversion * Remove dedicated dtype tests as per feedback --------- Co-authored-by: Dhruv Nair --- src/diffusers/pipelines/chroma/pipeline_chroma.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/diffusers/pipelines/chroma/pipeline_chroma.py b/src/diffusers/pipelines/chroma/pipeline_chroma.py index 19ea7729c9..5482035b3a 100644 --- a/src/diffusers/pipelines/chroma/pipeline_chroma.py +++ b/src/diffusers/pipelines/chroma/pipeline_chroma.py @@ -238,7 +238,7 @@ class ChromaPipeline( # Chroma requires the attention mask to include one padding token seq_lengths = attention_mask.sum(dim=1) mask_indices = torch.arange(attention_mask.size(1)).unsqueeze(0).expand(batch_size, -1) - attention_mask = (mask_indices <= seq_lengths.unsqueeze(1)).long() + attention_mask = (mask_indices <= seq_lengths.unsqueeze(1)).bool() prompt_embeds = self.text_encoder( text_input_ids.to(device), output_hidden_states=False, attention_mask=attention_mask.to(device) @@ -246,7 +246,7 @@ class ChromaPipeline( dtype = self.text_encoder.dtype prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) - attention_mask = attention_mask.to(dtype=dtype, device=device) + attention_mask = attention_mask.to(device=device) _, seq_len, _ = prompt_embeds.shape @@ -605,10 +605,9 @@ class ChromaPipeline( # Extend the prompt attention mask to account for image tokens in the final sequence attention_mask = torch.cat( - [attention_mask, torch.ones(batch_size, sequence_length, device=attention_mask.device)], + [attention_mask, torch.ones(batch_size, sequence_length, device=attention_mask.device, dtype=torch.bool)], dim=1, ) - attention_mask = attention_mask.to(dtype) return attention_mask From 64a5187d96f9376c7cf5123db810f2d2da79d7d0 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Mon, 29 Sep 2025 18:04:18 +0530 Subject: [PATCH 033/123] [quantization] feat: support aobaseconfig classes in `TorchAOConfig` (#12275) * feat: support aobaseconfig classes. * [docs] AOBaseConfig (#12302) init Co-authored-by: Sayak Paul * up * replace with is_torchao_version * up * up --------- Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> --- docs/source/en/quantization/torchao.md | 111 +++++++----- .../quantizers/quantization_config.py | 165 +++++++++++++++--- .../quantizers/torchao/torchao_quantizer.py | 90 +++++++--- tests/quantization/torchao/test_torchao.py | 22 +++ 4 files changed, 297 insertions(+), 91 deletions(-) diff --git a/docs/source/en/quantization/torchao.md b/docs/source/en/quantization/torchao.md index 5c7578dcbb..18cc109e07 100644 --- a/docs/source/en/quantization/torchao.md +++ b/docs/source/en/quantization/torchao.md @@ -11,69 +11,96 @@ specific language governing permissions and limitations under the License. --> # torchao -[TorchAO](https://github.com/pytorch/ao) is an architecture optimization library for PyTorch. It provides high-performance dtypes, optimization techniques, and kernels for inference and training, featuring composability with native PyTorch features like [torch.compile](https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html), FullyShardedDataParallel (FSDP), and more. +[torchao](https://github.com/pytorch/ao) provides high-performance dtypes and optimizations based on quantization and sparsity for inference and training PyTorch models. It is supported for any model in any modality, as long as it supports loading with [Accelerate](https://hf.co/docs/accelerate/index) and contains `torch.nn.Linear` layers. -Before you begin, make sure you have Pytorch 2.5+ and TorchAO installed. +Make sure Pytorch 2.5+ and torchao are installed with the command below. ```bash -pip install -U torch torchao +uv pip install -U torch torchao ``` +Each quantization dtype is available as a separate instance of a [AOBaseConfig](https://docs.pytorch.org/ao/main/api_ref_quantization.html#inference-apis-for-quantize) class. This provides more flexible configuration options by exposing more available arguments. -Quantize a model by passing [`TorchAoConfig`] to [`~ModelMixin.from_pretrained`] (you can also load pre-quantized models). This works for any model in any modality, as long as it supports loading with [Accelerate](https://hf.co/docs/accelerate/index) and contains `torch.nn.Linear` layers. +Pass the `AOBaseConfig` of a quantization dtype, like [Int4WeightOnlyConfig](https://docs.pytorch.org/ao/main/generated/torchao.quantization.Int4WeightOnlyConfig) to [`TorchAoConfig`] in [`~ModelMixin.from_pretrained`]. -The example below only quantizes the weights to int8. +```py +import torch +from diffusers import DiffusionPipeline, PipelineQuantizationConfig, TorchAoConfig +from torchao.quantization import Int8WeightOnlyConfig + +pipeline_quant_config = PipelineQuantizationConfig( + quant_mapping={"transformer": TorchAoConfig(Int8WeightOnlyConfig(group_size=128)))} +) +pipeline = DiffusionPipeline.from_pretrained( + "black-forest-labs/FLUX.1-dev", + quantzation_config=pipeline_quant_config, + torch_dtype=torch.bfloat16, + device_map="cuda" +) +``` + +For simple use cases, you could also provide a string identifier in [`TorchAo`] as shown below. + +```py +import torch +from diffusers import DiffusionPipeline, PipelineQuantizationConfig, TorchAoConfig + +pipeline_quant_config = PipelineQuantizationConfig( + quant_mapping={"transformer": TorchAoConfig("int8wo")} +) +pipeline = DiffusionPipeline.from_pretrained( + "black-forest-labs/FLUX.1-dev", + quantzation_config=pipeline_quant_config, + torch_dtype=torch.bfloat16, + device_map="cuda" +) +``` + +## torch.compile + +torchao supports [torch.compile](../optimization/fp16#torchcompile) which can speed up inference with one line of code. ```python import torch -from diffusers import FluxPipeline, AutoModel, TorchAoConfig +from diffusers import DiffusionPipeline, PipelineQuantizationConfig, TorchAoConfig +from torchao.quantization import Int4WeightOnlyConfig -model_id = "black-forest-labs/FLUX.1-dev" -dtype = torch.bfloat16 - -quantization_config = TorchAoConfig("int8wo") -transformer = AutoModel.from_pretrained( - model_id, - subfolder="transformer", - quantization_config=quantization_config, - torch_dtype=dtype, +pipeline_quant_config = PipelineQuantizationConfig( + quant_mapping={"transformer": TorchAoConfig(Int4WeightOnlyConfig(group_size=128)))} ) -pipe = FluxPipeline.from_pretrained( - model_id, - transformer=transformer, - torch_dtype=dtype, +pipeline = DiffusionPipeline.from_pretrained( + "black-forest-labs/FLUX.1-dev", + quantzation_config=pipeline_quant_config, + torch_dtype=torch.bfloat16, + device_map="cuda" ) -pipe.to("cuda") -# Without quantization: ~31.447 GB -# With quantization: ~20.40 GB -print(f"Pipeline memory usage: {torch.cuda.max_memory_reserved() / 1024**3:.3f} GB") - -prompt = "A cat holding a sign that says hello world" -image = pipe( - prompt, num_inference_steps=50, guidance_scale=4.5, max_sequence_length=512 -).images[0] -image.save("output.png") +pipeline.transformer.compile(transformer, mode="max-autotune", fullgraph=True) ``` -TorchAO is fully compatible with [torch.compile](../optimization/fp16#torchcompile), setting it apart from other quantization methods. This makes it easy to speed up inference with just one line of code. - -```python -# In the above code, add the following after initializing the transformer -transformer = torch.compile(transformer, mode="max-autotune", fullgraph=True) -``` - -For speed and memory benchmarks on Flux and CogVideoX, please refer to the table [here](https://github.com/huggingface/diffusers/pull/10009#issue-2688781450). You can also find some torchao [benchmarks](https://github.com/pytorch/ao/tree/main/torchao/quantization#benchmarks) numbers for various hardware. +Refer to this [table](https://github.com/huggingface/diffusers/pull/10009#issue-2688781450) for inference speed and memory usage benchmarks with Flux and CogVideoX. More benchmarks on various hardware are also available in the torchao [repository](https://github.com/pytorch/ao/tree/main/torchao/quantization#benchmarks). > [!TIP] > The FP8 post-training quantization schemes in torchao are effective for GPUs with compute capability of at least 8.9 (RTX-4090, Hopper, etc.). FP8 often provides the best speed, memory, and quality trade-off when generating images and videos. We recommend combining FP8 and torch.compile if your GPU is compatible. -torchao also supports an automatic quantization API through [autoquant](https://github.com/pytorch/ao/blob/main/torchao/quantization/README.md#autoquantization). Autoquantization determines the best quantization strategy applicable to a model by comparing the performance of each technique on chosen input types and shapes. Currently, this can be used directly on the underlying modeling components. Diffusers will also expose an autoquant configuration option in the future. +## autoquant -The `TorchAoConfig` class accepts three parameters: -- `quant_type`: A string value mentioning one of the quantization types below. -- `modules_to_not_convert`: A list of module full/partial module names for which quantization should not be performed. For example, to not perform any quantization of the [`FluxTransformer2DModel`]'s first block, one would specify: `modules_to_not_convert=["single_transformer_blocks.0"]`. -- `kwargs`: A dict of keyword arguments to pass to the underlying quantization method which will be invoked based on `quant_type`. +torchao provides [autoquant](https://docs.pytorch.org/ao/stable/generated/torchao.quantization.autoquant.html#torchao.quantization.autoquant) an automatic quantization API. Autoquantization chooses the best quantization strategy by comparing the performance of each strategy on chosen input types and shapes. This is only supported in Diffusers for individual models at the moment. + +```py +import torch +from diffusers import DiffusionPipeline +from torchao.quantization import autoquant + +# Load the pipeline +pipeline = DiffusionPipeline.from_pretrained( + "black-forest-labs/FLUX.1-schnell", + torch_dtype=torch.bfloat16, + device_map="cuda" +) + +transformer = autoquant(pipeline.transformer) +``` ## Supported quantization types diff --git a/src/diffusers/quantizers/quantization_config.py b/src/diffusers/quantizers/quantization_config.py index bf85795651..5dd8f56717 100644 --- a/src/diffusers/quantizers/quantization_config.py +++ b/src/diffusers/quantizers/quantization_config.py @@ -21,19 +21,20 @@ https://github.com/huggingface/transformers/blob/52cb4034ada381fe1ffe8d428a1076e """ import copy +import dataclasses import importlib.metadata import inspect import json import os import warnings -from dataclasses import dataclass +from dataclasses import dataclass, is_dataclass from enum import Enum from functools import partial from typing import Any, Callable, Dict, List, Optional, Union from packaging import version -from ..utils import is_torch_available, is_torchao_available, logging +from ..utils import is_torch_available, is_torchao_available, is_torchao_version, logging if is_torch_available(): @@ -443,7 +444,7 @@ class TorchAoConfig(QuantizationConfigMixin): """This is a config class for torchao quantization/sparsity techniques. Args: - quant_type (`str`): + quant_type (Union[`str`, AOBaseConfig]): The type of quantization we want to use, currently supporting: - **Integer quantization:** - Full function names: `int4_weight_only`, `int8_dynamic_activation_int4_weight`, @@ -465,6 +466,7 @@ class TorchAoConfig(QuantizationConfigMixin): - **Unsigned Integer quantization:** - Full function names: `uintx_weight_only` - Shorthands: `uint1wo`, `uint2wo`, `uint3wo`, `uint4wo`, `uint5wo`, `uint6wo`, `uint7wo` + - An AOBaseConfig instance: for more advanced configuration options. modules_to_not_convert (`List[str]`, *optional*, default to `None`): The list of modules to not quantize, useful for quantizing models that explicitly require to have some modules left in their original precision. @@ -478,6 +480,12 @@ class TorchAoConfig(QuantizationConfigMixin): ```python from diffusers import FluxTransformer2DModel, TorchAoConfig + # AOBaseConfig-based configuration + from torchao.quantization import Int8WeightOnlyConfig + + quantization_config = TorchAoConfig(Int8WeightOnlyConfig()) + + # String-based config quantization_config = TorchAoConfig("int8wo") transformer = FluxTransformer2DModel.from_pretrained( "black-forest-labs/Flux.1-Dev", @@ -490,7 +498,7 @@ class TorchAoConfig(QuantizationConfigMixin): def __init__( self, - quant_type: str, + quant_type: Union[str, "AOBaseConfig"], # noqa: F821 modules_to_not_convert: Optional[List[str]] = None, **kwargs, ) -> None: @@ -504,34 +512,103 @@ class TorchAoConfig(QuantizationConfigMixin): else: self.quant_type_kwargs = kwargs - TORCHAO_QUANT_TYPE_METHODS = self._get_torchao_quant_type_to_method() - if self.quant_type not in TORCHAO_QUANT_TYPE_METHODS.keys(): - is_floating_quant_type = self.quant_type.startswith("float") or self.quant_type.startswith("fp") - if is_floating_quant_type and not self._is_xpu_or_cuda_capability_atleast_8_9(): + self.post_init() + + def post_init(self): + if not isinstance(self.quant_type, str): + if is_torchao_version("<=", "0.9.0"): raise ValueError( - f"Requested quantization type: {self.quant_type} is not supported on GPUs with CUDA capability <= 8.9. You " - f"can check the CUDA capability of your GPU using `torch.cuda.get_device_capability()`." + f"torchao <= 0.9.0 only supports string quant_type, got {type(self.quant_type).__name__}. " + f"Upgrade to torchao > 0.9.0 to use AOBaseConfig." ) - raise ValueError( - f"Requested quantization type: {self.quant_type} is not supported or is an incorrect `quant_type` name. If you think the " - f"provided quantization type should be supported, please open an issue at https://github.com/huggingface/diffusers/issues." - ) + from torchao.quantization.quant_api import AOBaseConfig - method = TORCHAO_QUANT_TYPE_METHODS[self.quant_type] - signature = inspect.signature(method) - all_kwargs = { - param.name - for param in signature.parameters.values() - if param.kind in [inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD] - } - unsupported_kwargs = list(self.quant_type_kwargs.keys() - all_kwargs) + if not isinstance(self.quant_type, AOBaseConfig): + raise TypeError(f"quant_type must be a AOBaseConfig instance, got {type(self.quant_type).__name__}") - if len(unsupported_kwargs) > 0: - raise ValueError( - f'The quantization method "{quant_type}" does not support the following keyword arguments: ' - f"{unsupported_kwargs}. The following keywords arguments are supported: {all_kwargs}." - ) + elif isinstance(self.quant_type, str): + TORCHAO_QUANT_TYPE_METHODS = self._get_torchao_quant_type_to_method() + + if self.quant_type not in TORCHAO_QUANT_TYPE_METHODS.keys(): + is_floating_quant_type = self.quant_type.startswith("float") or self.quant_type.startswith("fp") + if is_floating_quant_type and not self._is_xpu_or_cuda_capability_atleast_8_9(): + raise ValueError( + f"Requested quantization type: {self.quant_type} is not supported on GPUs with CUDA capability <= 8.9. You " + f"can check the CUDA capability of your GPU using `torch.cuda.get_device_capability()`." + ) + + raise ValueError( + f"Requested quantization type: {self.quant_type} is not supported or is an incorrect `quant_type` name. If you think the " + f"provided quantization type should be supported, please open an issue at https://github.com/huggingface/diffusers/issues." + ) + + method = TORCHAO_QUANT_TYPE_METHODS[self.quant_type] + signature = inspect.signature(method) + all_kwargs = { + param.name + for param in signature.parameters.values() + if param.kind in [inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD] + } + unsupported_kwargs = list(self.quant_type_kwargs.keys() - all_kwargs) + + if len(unsupported_kwargs) > 0: + raise ValueError( + f'The quantization method "{self.quant_type}" does not support the following keyword arguments: ' + f"{unsupported_kwargs}. The following keywords arguments are supported: {all_kwargs}." + ) + + def to_dict(self): + """Convert configuration to a dictionary.""" + d = super().to_dict() + + if isinstance(self.quant_type, str): + # Handle layout serialization if present + if "quant_type_kwargs" in d and "layout" in d["quant_type_kwargs"]: + if is_dataclass(d["quant_type_kwargs"]["layout"]): + d["quant_type_kwargs"]["layout"] = [ + d["quant_type_kwargs"]["layout"].__class__.__name__, + dataclasses.asdict(d["quant_type_kwargs"]["layout"]), + ] + if isinstance(d["quant_type_kwargs"]["layout"], list): + assert len(d["quant_type_kwargs"]["layout"]) == 2, "layout saves layout name and layout kwargs" + assert isinstance(d["quant_type_kwargs"]["layout"][0], str), "layout name must be a string" + assert isinstance(d["quant_type_kwargs"]["layout"][1], dict), "layout kwargs must be a dict" + else: + raise ValueError("layout must be a list") + else: + # Handle AOBaseConfig serialization + from torchao.core.config import config_to_dict + + # For now we assume there is 1 config per Transformer, however in the future + # We may want to support a config per fqn. + d["quant_type"] = {"default": config_to_dict(self.quant_type)} + + return d + + @classmethod + def from_dict(cls, config_dict, return_unused_kwargs=False, **kwargs): + """Create configuration from a dictionary.""" + if not is_torchao_version(">", "0.9.0"): + raise NotImplementedError("TorchAoConfig requires torchao > 0.9.0 for construction from dict") + config_dict = config_dict.copy() + quant_type = config_dict.pop("quant_type") + + if isinstance(quant_type, str): + return cls(quant_type=quant_type, **config_dict) + # Check if we only have one key which is "default" + # In the future we may update this + assert len(quant_type) == 1 and "default" in quant_type, ( + "Expected only one key 'default' in quant_type dictionary" + ) + quant_type = quant_type["default"] + + # Deserialize quant_type if needed + from torchao.core.config import config_from_dict + + quant_type = config_from_dict(quant_type) + + return cls(quant_type=quant_type, **config_dict) @classmethod def _get_torchao_quant_type_to_method(cls): @@ -681,8 +758,38 @@ class TorchAoConfig(QuantizationConfigMixin): raise RuntimeError("TorchAO requires a CUDA compatible GPU or Intel XPU and installation of PyTorch.") def get_apply_tensor_subclass(self): - TORCHAO_QUANT_TYPE_METHODS = self._get_torchao_quant_type_to_method() - return TORCHAO_QUANT_TYPE_METHODS[self.quant_type](**self.quant_type_kwargs) + """Create the appropriate quantization method based on configuration.""" + if not isinstance(self.quant_type, str): + return self.quant_type + else: + methods = self._get_torchao_quant_type_to_method() + quant_type_kwargs = self.quant_type_kwargs.copy() + if ( + not torch.cuda.is_available() + and is_torchao_available() + and self.quant_type == "int4_weight_only" + and version.parse(importlib.metadata.version("torchao")) >= version.parse("0.8.0") + and quant_type_kwargs.get("layout", None) is None + ): + if torch.xpu.is_available(): + if version.parse(importlib.metadata.version("torchao")) >= version.parse( + "0.11.0" + ) and version.parse(importlib.metadata.version("torch")) > version.parse("2.7.9"): + from torchao.dtypes import Int4XPULayout + from torchao.quantization.quant_primitives import ZeroPointDomain + + quant_type_kwargs["layout"] = Int4XPULayout() + quant_type_kwargs["zero_point_domain"] = ZeroPointDomain.INT + else: + raise ValueError( + "TorchAoConfig requires torchao >= 0.11.0 and torch >= 2.8.0 for XPU support. Please upgrade the version or use run on CPU with the cpu version pytorch." + ) + else: + from torchao.dtypes import Int4CPULayout + + quant_type_kwargs["layout"] = Int4CPULayout() + + return methods[self.quant_type](**quant_type_kwargs) def __repr__(self): r""" diff --git a/src/diffusers/quantizers/torchao/torchao_quantizer.py b/src/diffusers/quantizers/torchao/torchao_quantizer.py index 976bc8a1e0..2334c7af86 100644 --- a/src/diffusers/quantizers/torchao/torchao_quantizer.py +++ b/src/diffusers/quantizers/torchao/torchao_quantizer.py @@ -18,9 +18,10 @@ https://github.com/huggingface/transformers/blob/3a8eb74668e9c2cc563b2f5c62fac17 """ import importlib +import re import types from fnmatch import fnmatch -from typing import TYPE_CHECKING, Any, Dict, List, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union from packaging import version @@ -107,6 +108,21 @@ if ( _update_torch_safe_globals() +def fuzzy_match_size(config_name: str) -> Optional[str]: + """ + Extract the size digit from strings like "4weight", "8weight". Returns the digit as an integer if found, otherwise + None. + """ + config_name = config_name.lower() + + str_match = re.search(r"(\d)weight", config_name) + + if str_match: + return str_match.group(1) + + return None + + logger = logging.get_logger(__name__) @@ -176,8 +192,7 @@ class TorchAoHfQuantizer(DiffusersQuantizer): def update_torch_dtype(self, torch_dtype): quant_type = self.quantization_config.quant_type - - if quant_type.startswith("int") or quant_type.startswith("uint"): + if isinstance(quant_type, str) and (quant_type.startswith("int") or quant_type.startswith("uint")): if torch_dtype is not None and torch_dtype != torch.bfloat16: logger.warning( f"You are trying to set torch_dtype to {torch_dtype} for int4/int8/uintx quantization, but " @@ -197,24 +212,44 @@ class TorchAoHfQuantizer(DiffusersQuantizer): def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype": quant_type = self.quantization_config.quant_type + from accelerate.utils import CustomDtype - if quant_type.startswith("int8") or quant_type.startswith("int4"): - # Note that int4 weights are created by packing into torch.int8, but since there is no torch.int4, we use torch.int8 - return torch.int8 - elif quant_type == "uintx_weight_only": - return self.quantization_config.quant_type_kwargs.get("dtype", torch.uint8) - elif quant_type.startswith("uint"): - return { - 1: torch.uint1, - 2: torch.uint2, - 3: torch.uint3, - 4: torch.uint4, - 5: torch.uint5, - 6: torch.uint6, - 7: torch.uint7, - }[int(quant_type[4])] - elif quant_type.startswith("float") or quant_type.startswith("fp"): - return torch.bfloat16 + if isinstance(quant_type, str): + if quant_type.startswith("int8"): + # Note that int4 weights are created by packing into torch.int8, but since there is no torch.int4, we use torch.int8 + return torch.int8 + elif quant_type.startswith("int4"): + return CustomDtype.INT4 + elif quant_type == "uintx_weight_only": + return self.quantization_config.quant_type_kwargs.get("dtype", torch.uint8) + elif quant_type.startswith("uint"): + return { + 1: torch.uint1, + 2: torch.uint2, + 3: torch.uint3, + 4: torch.uint4, + 5: torch.uint5, + 6: torch.uint6, + 7: torch.uint7, + }[int(quant_type[4])] + elif quant_type.startswith("float") or quant_type.startswith("fp"): + return torch.bfloat16 + + elif is_torchao_version(">", "0.9.0"): + from torchao.core.config import AOBaseConfig + + quant_type = self.quantization_config.quant_type + if isinstance(quant_type, AOBaseConfig): + # Extract size digit using fuzzy match on the class name + config_name = quant_type.__class__.__name__ + size_digit = fuzzy_match_size(config_name) + + # Map the extracted digit to appropriate dtype + if size_digit == "4": + return CustomDtype.INT4 + else: + # Default to int8 + return torch.int8 if isinstance(target_dtype, SUPPORTED_TORCH_DTYPES_FOR_QUANTIZATION): return target_dtype @@ -297,6 +332,21 @@ class TorchAoHfQuantizer(DiffusersQuantizer): # Original mapping for non-AOBaseConfig types # For the uint types, this is a best guess. Once these types become more used # we can look into their nuances. + if is_torchao_version(">", "0.9.0"): + from torchao.core.config import AOBaseConfig + + quant_type = self.quantization_config.quant_type + # For autoquant case, it will be treated in the string implementation below in map_to_target_dtype + if isinstance(quant_type, AOBaseConfig): + # Extract size digit using fuzzy match on the class name + config_name = quant_type.__class__.__name__ + size_digit = fuzzy_match_size(config_name) + + if size_digit == "4": + return 8 + else: + return 4 + map_to_target_dtype = {"int4_*": 8, "int8_*": 4, "uint*": 8, "float8*": 4} quant_type = self.quantization_config.quant_type for pattern, target_dtype in map_to_target_dtype.items(): diff --git a/tests/quantization/torchao/test_torchao.py b/tests/quantization/torchao/test_torchao.py index 920c3a55f5..38997de17b 100644 --- a/tests/quantization/torchao/test_torchao.py +++ b/tests/quantization/torchao/test_torchao.py @@ -14,11 +14,13 @@ # limitations under the License. import gc +import importlib.metadata import tempfile import unittest from typing import List import numpy as np +from packaging import version from parameterized import parameterized from transformers import AutoTokenizer, CLIPTextModel, CLIPTokenizer, T5EncoderModel @@ -65,6 +67,9 @@ if is_torchao_available(): from torchao.quantization.quant_primitives import MappingType from torchao.utils import get_model_size_in_bytes + if version.parse(importlib.metadata.version("torchao")) >= version.Version("0.9.0"): + from torchao.quantization import Int8WeightOnlyConfig + @require_torch @require_torch_accelerator @@ -522,6 +527,15 @@ class TorchAoTest(unittest.TestCase): inputs = self.get_dummy_inputs(torch_device) _ = pipe(**inputs) + @require_torchao_version_greater_or_equal("0.9.0") + def test_aobase_config(self): + quantization_config = TorchAoConfig(Int8WeightOnlyConfig()) + components = self.get_dummy_components(quantization_config) + pipe = FluxPipeline(**components).to(torch_device) + + inputs = self.get_dummy_inputs(torch_device) + _ = pipe(**inputs) + # Slices for these tests have been obtained on our aws-g6e-xlarge-plus runners @require_torch @@ -628,6 +642,14 @@ class TorchAoSerializationTest(unittest.TestCase): self._test_original_model_expected_slice(quant_method, quant_method_kwargs, expected_slice) self._check_serialization_expected_slice(quant_method, quant_method_kwargs, expected_slice, device) + @require_torchao_version_greater_or_equal("0.9.0") + def test_aobase_config(self): + quant_method, quant_method_kwargs = Int8WeightOnlyConfig(), {} + expected_slice = np.array([0.3613, -0.127, -0.0223, -0.2539, -0.459, 0.4961, -0.1357, -0.6992, 0.4551]) + device = torch_device + self._test_original_model_expected_slice(quant_method, quant_method_kwargs, expected_slice) + self._check_serialization_expected_slice(quant_method, quant_method_kwargs, expected_slice, device) + @require_torchao_version_greater_or_equal("0.7.0") class TorchAoCompileTest(QuantCompileTests, unittest.TestCase): From ccedeca96e9aebd3e0e663e668891bea9de30dbc Mon Sep 17 00:00:00 2001 From: Steven Liu <59462357+stevhliu@users.noreply.github.com> Date: Mon, 29 Sep 2025 11:24:26 -0700 Subject: [PATCH 034/123] [docs] Distributed inference (#12285) * init * feedback --------- Co-authored-by: Sayak Paul --- .../en/training/distributed_inference.md | 96 +++++++++---------- 1 file changed, 45 insertions(+), 51 deletions(-) diff --git a/docs/source/en/training/distributed_inference.md b/docs/source/en/training/distributed_inference.md index a536703f5b..58ec77f75b 100644 --- a/docs/source/en/training/distributed_inference.md +++ b/docs/source/en/training/distributed_inference.md @@ -12,17 +12,23 @@ specific language governing permissions and limitations under the License. # Distributed inference -On distributed setups, you can run inference across multiple GPUs with 🤗 [Accelerate](https://huggingface.co/docs/accelerate/index) or [PyTorch Distributed](https://pytorch.org/tutorials/beginner/dist_overview.html), which is useful for generating with multiple prompts in parallel. +Distributed inference splits the workload across multiple GPUs. It a useful technique for fitting larger models in memory and can process multiple prompts for higher throughput. -This guide will show you how to use 🤗 Accelerate and PyTorch Distributed for distributed inference. +This guide will show you how to use [Accelerate](https://huggingface.co/docs/accelerate/index) and [PyTorch Distributed](https://pytorch.org/tutorials/beginner/dist_overview.html) for distributed inference. -## 🤗 Accelerate +## Accelerate -🤗 [Accelerate](https://huggingface.co/docs/accelerate/index) is a library designed to make it easy to train or run inference across distributed setups. It simplifies the process of setting up the distributed environment, allowing you to focus on your PyTorch code. +Accelerate is a library designed to simplify inference and training on multiple accelerators by handling the setup, allowing users to focus on their PyTorch code. -To begin, create a Python file and initialize an [`accelerate.PartialState`] to create a distributed environment; your setup is automatically detected so you don't need to explicitly define the `rank` or `world_size`. Move the [`DiffusionPipeline`] to `distributed_state.device` to assign a GPU to each process. +Install Accelerate with the following command. -Now use the [`~accelerate.PartialState.split_between_processes`] utility as a context manager to automatically distribute the prompts between the number of processes. +```bash +uv pip install accelerate +``` + +Initialize a [`accelerate.PartialState`] class in a Python file to create a distributed environment. The [`accelerate.PartialState`] class manages process management, device control and distribution, and process coordination. + +Move the [`DiffusionPipeline`] to [`accelerate.PartialState.device`] to assign a GPU to each process. ```py import torch @@ -30,33 +36,31 @@ from accelerate import PartialState from diffusers import DiffusionPipeline pipeline = DiffusionPipeline.from_pretrained( - "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True + "Qwen/Qwen-Image", torch_dtype=torch.float16 ) distributed_state = PartialState() pipeline.to(distributed_state.device) +``` +Use the [`~accelerate.PartialState.split_between_processes`] utility as a context manager to automatically distribute the prompts between the number of processes. + +```py with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt: result = pipeline(prompt).images[0] result.save(f"result_{distributed_state.process_index}.png") ``` -Use the `--num_processes` argument to specify the number of GPUs to use, and call `accelerate launch` to run the script: +Call `accelerate launch` to run the script and use the `--num_processes` argument to set the number of GPUs to use. ```bash accelerate launch run_distributed.py --num_processes=2 ``` - - -Refer to this minimal example [script](https://gist.github.com/sayakpaul/cfaebd221820d7b43fae638b4dfa01ba) for running inference across multiple GPUs. To learn more, take a look at the [Distributed Inference with 🤗 Accelerate](https://huggingface.co/docs/accelerate/en/usage_guides/distributed_inference#distributed-inference-with-accelerate) guide. - - - ## PyTorch Distributed -PyTorch supports [`DistributedDataParallel`](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html) which enables data parallelism. +PyTorch [DistributedDataParallel](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html) enables [data parallelism](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=data_parallelism), which replicates the same model on each device, to process different batches of data in parallel. -To start, create a Python file and import `torch.distributed` and `torch.multiprocessing` to set up the distributed process group and to spawn the processes for inference on each GPU. You should also initialize a [`DiffusionPipeline`]: +Import `torch.distributed` and `torch.multiprocessing` into a Python file to set up the distributed process group and to spawn the processes for inference on each GPU. ```py import torch @@ -65,20 +69,20 @@ import torch.multiprocessing as mp from diffusers import DiffusionPipeline -sd = DiffusionPipeline.from_pretrained( - "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True +pipeline = DiffusionPipeline.from_pretrained( + "Qwen/Qwen-Image", torch_dtype=torch.float16, ) ``` -You'll want to create a function to run inference; [`init_process_group`](https://pytorch.org/docs/stable/distributed.html?highlight=init_process_group#torch.distributed.init_process_group) handles creating a distributed environment with the type of backend to use, the `rank` of the current process, and the `world_size` or the number of processes participating. If you're running inference in parallel over 2 GPUs, then the `world_size` is 2. +Create a function for inference with [init_process_group](https://pytorch.org/docs/stable/distributed.html?highlight=init_process_group#torch.distributed.init_process_group). This method creates a distributed environment with the backend type, the `rank` of the current process, and the `world_size` or number of processes participating (for example, 2 GPUs would be `world_size=2`). -Move the [`DiffusionPipeline`] to `rank` and use `get_rank` to assign a GPU to each process, where each process handles a different prompt: +Move the pipeline to `rank` and use `get_rank` to assign a GPU to each process. Each process handles a different prompt. ```py def run_inference(rank, world_size): dist.init_process_group("nccl", rank=rank, world_size=world_size) - sd.to(rank) + pipeline.to(rank) if torch.distributed.get_rank() == 0: prompt = "a dog" @@ -89,7 +93,7 @@ def run_inference(rank, world_size): image.save(f"./{'_'.join(prompt)}.png") ``` -To run the distributed inference, call [`mp.spawn`](https://pytorch.org/docs/stable/multiprocessing.html#torch.multiprocessing.spawn) to run the `run_inference` function on the number of GPUs defined in `world_size`: +Use [mp.spawn](https://pytorch.org/docs/stable/multiprocessing.html#torch.multiprocessing.spawn) to create the number of processes defined in `world_size`. ```py def main(): @@ -101,31 +105,26 @@ if __name__ == "__main__": main() ``` -Once you've completed the inference script, use the `--nproc_per_node` argument to specify the number of GPUs to use and call `torchrun` to run the script: +Call `torchrun` to run the inference script and use the `--nproc_per_node` argument to set the number of GPUs to use. ```bash torchrun run_distributed.py --nproc_per_node=2 ``` -> [!TIP] -> You can use `device_map` within a [`DiffusionPipeline`] to distribute its model-level components on multiple devices. Refer to the [Device placement](../tutorials/inference_with_big_models#device-placement) guide to learn more. +## device_map -## Model sharding +The `device_map` argument enables distributed inference by automatically placing model components on separate GPUs. This is especially useful when a model doesn't fit on a single GPU. You can use `device_map` to selectively load and unload the required model components at a given stage as shown in the example below (assumes two GPUs are available). -Modern diffusion systems such as [Flux](../api/pipelines/flux) are very large and have multiple models. For example, [Flux.1-Dev](https://hf.co/black-forest-labs/FLUX.1-dev) is made up of two text encoders - [T5-XXL](https://hf.co/google/t5-v1_1-xxl) and [CLIP-L](https://hf.co/openai/clip-vit-large-patch14) - a [diffusion transformer](../api/models/flux_transformer), and a [VAE](../api/models/autoencoderkl). With a model this size, it can be challenging to run inference on consumer GPUs. - -Model sharding is a technique that distributes models across GPUs when the models don't fit on a single GPU. The example below assumes two 16GB GPUs are available for inference. - -Start by computing the text embeddings with the text encoders. Keep the text encoders on two GPUs by setting `device_map="balanced"`. The `balanced` strategy evenly distributes the model on all available GPUs. Use the `max_memory` parameter to allocate the maximum amount of memory for each text encoder on each GPU. - -> [!TIP] -> **Only** load the text encoders for this step! The diffusion transformer and VAE are loaded in a later step to preserve memory. +Set `device_map="balanced"` to evenly distributes the text encoders on all available GPUs. You can use the `max_memory` argument to allocate a maximum amount of memory for each text encoder. Don't load any other pipeline components to avoid memory usage. ```py from diffusers import FluxPipeline import torch -prompt = "a photo of a dog with cat-like look" +prompt = """ +cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California +highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain +""" pipeline = FluxPipeline.from_pretrained( "black-forest-labs/FLUX.1-dev", @@ -142,7 +141,7 @@ with torch.no_grad(): ) ``` -Once the text embeddings are computed, remove them from the GPU to make space for the diffusion transformer. +After the text embeddings are computed, remove them from the GPU to make space for the diffusion transformer. ```py import gc @@ -162,7 +161,7 @@ del pipeline flush() ``` -Load the diffusion transformer next which has 12.5B parameters. This time, set `device_map="auto"` to automatically distribute the model across two 16GB GPUs. The `auto` strategy is backed by [Accelerate](https://hf.co/docs/accelerate/index) and available as a part of the [Big Model Inference](https://hf.co/docs/accelerate/concept_guides/big_model_inference) feature. It starts by distributing a model across the fastest device first (GPU) before moving to slower devices like the CPU and hard drive if needed. The trade-off of storing model parameters on slower devices is slower inference latency. +Set `device_map="auto"` to automatically distribute the model on the two GPUs. This strategy places a model on the fastest device first before placing a model on a slower device like a CPU or hard drive if needed. The trade-off of storing model parameters on slower devices is slower inference latency. ```py from diffusers import AutoModel @@ -177,9 +176,9 @@ transformer = AutoModel.from_pretrained( ``` > [!TIP] -> At any point, you can try `print(pipeline.hf_device_map)` to see how the various models are distributed across devices. This is useful for tracking the device placement of the models. You can also try `print(transformer.hf_device_map)` to see how the transformer model is sharded across devices. +> Run `pipeline.hf_device_map` to see how the various models are distributed across devices. This is useful for tracking model device placement. You can also call `hf_device_map` on the transformer model to see how it is distributed. -Add the transformer model to the pipeline for denoising, but set the other model-level components like the text encoders and VAE to `None` because you don't need them yet. +Add the transformer model to the pipeline and set the `output_type="latent"` to generate the latents. ```py pipeline = FluxPipeline.from_pretrained( @@ -206,21 +205,12 @@ latents = pipeline( ).images ``` -Remove the pipeline and transformer from memory as they're no longer needed. - -```py -del pipeline.transformer -del pipeline - -flush() -``` - -Finally, decode the latents with the VAE into an image. The VAE is typically small enough to be loaded on a single GPU. +Remove the pipeline and transformer from memory and load a VAE to decode the latents. The VAE is typically small enough to be loaded on a single device. ```py +import torch from diffusers import AutoencoderKL from diffusers.image_processor import VaeImageProcessor -import torch vae = AutoencoderKL.from_pretrained(ckpt_id, subfolder="vae", torch_dtype=torch.bfloat16).to("cuda") vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1) @@ -236,4 +226,8 @@ with torch.no_grad(): image[0].save("split_transformer.png") ``` -By selectively loading and unloading the models you need at a given stage and sharding the largest models across multiple GPUs, it is possible to run inference with large models on consumer GPUs. +## Resources + +- Take a look at this [script](https://gist.github.com/sayakpaul/cfaebd221820d7b43fae638b4dfa01ba) for a minimal example of distributed inference with Accelerate. +- For more details, check out Accelerate's [Distributed inference](https://huggingface.co/docs/accelerate/en/usage_guides/distributed_inference#distributed-inference-with-accelerate) guide. +- The `device_map` argument assign models or an entire pipeline to devices. Refer to the [device placement](../using-diffusers/loading#device-placement) docs for more information. \ No newline at end of file From c07fcf780a199be58018a80349e8447077146ac5 Mon Sep 17 00:00:00 2001 From: Steven Liu <59462357+stevhliu@users.noreply.github.com> Date: Mon, 29 Sep 2025 11:36:14 -0700 Subject: [PATCH 035/123] [docs] Model formats (#12256) * init * config * lora metadata * feedback * fix * cache allocator warmup for from_single_file * feedback * feedback --- docs/source/en/_toctree.yml | 2 +- docs/source/en/using-diffusers/loading.md | 3 + .../en/using-diffusers/other-formats.md | 607 ++++++------------ 3 files changed, 193 insertions(+), 419 deletions(-) diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 96c6fbb17f..ada5e38895 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -25,7 +25,7 @@ - local: using-diffusers/schedulers title: Schedulers - local: using-diffusers/other-formats - title: Model files and layouts + title: Model formats - local: using-diffusers/push_to_hub title: Sharing pipelines and models diff --git a/docs/source/en/using-diffusers/loading.md b/docs/source/en/using-diffusers/loading.md index 25b53d2f4d..3fb608b1c2 100644 --- a/docs/source/en/using-diffusers/loading.md +++ b/docs/source/en/using-diffusers/loading.md @@ -52,6 +52,9 @@ pipeline = QwenImagePipeline.from_pretrained( ) ``` +> [!TIP] +> Refer to the [Single file format](./other-formats#single-file-format) docs to learn how to load single file models. + ### Local pipelines Pipelines can also be run locally. Use [`~huggingface_hub.snapshot_download`] to download a model repository. diff --git a/docs/source/en/using-diffusers/other-formats.md b/docs/source/en/using-diffusers/other-formats.md index 59835bbf26..b6e333ed77 100644 --- a/docs/source/en/using-diffusers/other-formats.md +++ b/docs/source/en/using-diffusers/other-formats.md @@ -10,77 +10,183 @@ an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express o specific language governing permissions and limitations under the License. --> -# Model files and layouts - [[open-in-colab]] -Diffusion models are saved in various file types and organized in different layouts. Diffusers stores model weights as safetensors files in *Diffusers-multifolder* layout and it also supports loading files (like safetensors and ckpt files) from a *single-file* layout which is commonly used in the diffusion ecosystem. +# Model formats -Each layout has its own benefits and use cases, and this guide will show you how to load the different files and layouts, and how to convert them. +Diffusion models are typically stored in the Diffusers format or single-file format. Model files can be stored in various file types such as safetensors, dduf, or ckpt. -## Files +> [!TIP] +> Format refers to whether the weights are stored in a directory structure and file refers to the file type. -PyTorch model weights are typically saved with Python's [pickle](https://docs.python.org/3/library/pickle.html) utility as ckpt or bin files. However, pickle is not secure and pickled files may contain malicious code that can be executed. This vulnerability is a serious concern given the popularity of model sharing. To address this security issue, the [Safetensors](https://hf.co/docs/safetensors) library was developed as a secure alternative to pickle, which saves models as safetensors files. +This guide will show you how to load pipelines and models from these formats and files. + +## Diffusers format + +The Diffusers format stores each model (UNet, transformer, text encoder) in a separate subfolder. There are several benefits to storing models separately. + +- Faster overall pipeline initialization because you can load the individual model you need or load them all in parallel. +- Reduced memory usage because you don't need to load all the pipeline components if you only need one model. [Reuse](./loading#reusing-models-in-multiple-pipelines) a model that is shared between multiple pipelines. +- Lower storage requirements because common models shared between multiple pipelines are only downloaded once. +- Flexibility to use new or improved models in a pipeline. + +## Single file format + +A single-file format stores *all* the model (UNet, transformer, text encoder) weights in a single file. Benefits of single-file formats include the following. + +- Greater compatibility with [ComfyUI](https://github.com/comfyanonymous/ComfyUI) or [Automatic1111](https://github.com/AUTOMATIC1111/stable-diffusion-webui). +- Easier to download and share a single file. + +Use [`~loaders.FromSingleFileMixin.from_single_file`] to load a single file. + +```py +import torch +from diffusers import StableDiffusionXLPipeline + +pipeline = StableDiffusionXLPipeline.from_single_file( + "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_base_1.0.safetensors", + torch_dtype=torch.float16, + device_map="cuda" +) +``` + +The [`~loaders.FromSingleFileMixin.from_single_file`] method also supports passing new models or schedulers. + +```py +import torch +from diffusers import FluxPipeline, FluxTransformer2DModel + +transformer = FluxTransformer2DModel.from_single_file( + "https://huggingface.co/Kijai/flux-fp8/blob/main/flux1-dev-fp8.safetensors", torch_dtype=torch.bfloat16 +) +pipeline = FluxPipeline.from_pretrained( + "black-forest-labs/FLUX.1-dev", + transformer=transformer, + torch_dtype=torch.bfloat16, + device_map="cuda" +) +``` + +### Configuration options + +Diffusers format models have a `config.json` file in their repositories with important attributes such as the number of layers and attention heads. The [`~loaders.FromSingleFileMixin.from_single_file`] method automatically determines the appropriate config to use from `config.json`. This may fail in a few rare instances though, in which case, you should use the `config` argument. + +You should also use the `config` argument if the models in a pipeline are different from the original implementation or if it doesn't have the necessary metadata to determine the correct config. + +```py +from diffusers import StableDiffusionXLPipeline + +ckpt_path = "https://huggingface.co/segmind/SSD-1B/blob/main/SSD-1B.safetensors" + +pipeline = StableDiffusionXLPipeline.from_single_file(ckpt_path, config="segmind/SSD-1B") +``` + +Diffusers attempts to infer the pipeline components based on the signature types of the pipeline class when using `original_config` with `local_files_only=True`. It won't download the config files from a Hub repository to avoid backward breaking changes when you can't connect to the internet. This method isn't as reliable as providing a path to a local model with the `config` argument and may lead to errors. You should run the pipeline with `local_files_only=False` to download the config files to the local cache to avoid errors. + +Override default configs by passing the arguments directly to [`~loaders.FromSingleFileMixin.from_single_file`]. The examples below demonstrate how to override the configs in a pipeline or model. + +```py +from diffusers import StableDiffusionXLInstructPix2PixPipeline + +ckpt_path = "https://huggingface.co/stabilityai/cosxl/blob/main/cosxl_edit.safetensors" +pipeline = StableDiffusionXLInstructPix2PixPipeline.from_single_file( + ckpt_path, config="diffusers/sdxl-instructpix2pix-768", is_cosxl_edit=True +) +``` + +```py +from diffusers import UNet2DConditionModel + +ckpt_path = "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_base_1.0_0.9vae.safetensors" +model = UNet2DConditionModel.from_single_file(ckpt_path, upcast_attention=True) +``` + +### Local files + +The [`~loaders.FromSingleFileMixin.from_single_file`] method attempts to configure a pipeline or model by inferring the model type from the keys in the checkpoint file. For example, any single file checkpoint based on the Stable Diffusion XL base model is configured from [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0). + +If you're working with local files, download the config files with the [`~huggingface_hub.snapshot_download`] method and the model checkpoint with [`~huggingface_hub.hf_hub_download`]. These files are downloaded to your [cache directory](https://huggingface.co/docs/huggingface_hub/en/guides/manage-cache), but you can download them to a specific directory with the `local_dir` argument. + +```py +from huggingface_hub import hf_hub_download, snapshot_download +from diffusers import StableDiffusionXLPipeline + +my_local_checkpoint_path = hf_hub_download( + repo_id="segmind/SSD-1B", + filename="SSD-1B.safetensors" +) + +my_local_config_path = snapshot_download( + repo_id="segmind/SSD-1B", + allow_patterns=["*.json", "**/*.json", "*.txt", "**/*.txt"] +) + +pipeline = StableDiffusionXLPipeline.from_single_file( + my_local_checkpoint_path, config=my_local_config_path, local_files_only=True +) +``` + +### Symlink + +If you're working with a file system that does not support symlinking, download the checkpoint file to a local directory first with the `local_dir` parameter. Using the `local_dir` parameter automatically disables symlinks. + +```py +from huggingface_hub import hf_hub_download, snapshot_download +from diffusers import StableDiffusionXLPipeline + +my_local_checkpoint_path = hf_hub_download( + repo_id="segmind/SSD-1B", + filename="SSD-1B.safetensors" + local_dir="my_local_checkpoints", +) +print("My local checkpoint: ", my_local_checkpoint_path) + +my_local_config_path = snapshot_download( + repo_id="segmind/SSD-1B", + allow_patterns=["*.json", "**/*.json", "*.txt", "**/*.txt"] +) +print("My local config: ", my_local_config_path) +``` + +Pass these paths to [`~loaders.FromSingleFileMixin.from_single_file`]. + +```py +pipeline = StableDiffusionXLPipeline.from_single_file( + my_local_checkpoint_path, config=my_local_config_path, local_files_only=True +) +``` + +## File types + +Models can be stored in several file types. Safetensors is the most common file type but you may encounter other file types on the Hub or diffusion community. ### safetensors -> [!TIP] -> Learn more about the design decisions and why safetensor files are preferred for saving and loading model weights in the [Safetensors audited as really safe and becoming the default](https://blog.eleuther.ai/safetensors-security-audit/) blog post. +[Safetensors](https://hf.co/docs/safetensors) is a safe and fast file type for securely storing and loading tensors. It restricts the header size to limit certain types of attacks, supports lazy loading (useful for distributed setups), and generally loads faster. -[Safetensors](https://hf.co/docs/safetensors) is a safe and fast file format for securely storing and loading tensors. Safetensors restricts the header size to limit certain types of attacks, supports lazy loading (useful for distributed setups), and has generally faster loading speeds. +Diffusers loads safetensors file by default (a required dependency) if they are available and the Safetensors library is installed. -Make sure you have the [Safetensors](https://hf.co/docs/safetensors) library installed. - -```py -!pip install safetensors -``` - -Safetensors stores weights in a safetensors file. Diffusers loads safetensors files by default if they're available and the Safetensors library is installed. There are two ways safetensors files can be organized: - -1. Diffusers-multifolder layout: there may be several separate safetensors files, one for each pipeline component (text encoder, UNet, VAE), organized in subfolders (check out the [stable-diffusion-v1-5/stable-diffusion-v1-5](https://hf.co/stable-diffusion-v1-5/stable-diffusion-v1-5/tree/main) repository as an example) -2. single-file layout: all the model weights may be saved in a single file (check out the [WarriorMama777/OrangeMixs](https://hf.co/WarriorMama777/OrangeMixs/tree/main/Models/AbyssOrangeMix) repository as an example) - - - - -Use the [`~DiffusionPipeline.from_pretrained`] method to load a model with safetensors files stored in multiple folders. +Use [`~DiffusionPipeline.from_pretrained`] or [`~loaders.FromSingleFileMixin.from_single_file`] to load safetensor files. ```py +import torch from diffusers import DiffusionPipeline pipeline = DiffusionPipeline.from_pretrained( - "stable-diffusion-v1-5/stable-diffusion-v1-5", - use_safetensors=True + "stabilityai/stable-diffusion-xl-base-1.0", + torch.dtype=torch.float16, + device_map="cuda" +) + +pipeline = DiffusionPipeline.from_single_file( + "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_base_1.0.safetensors", + torch_dtype=torch.float16, ) ``` - - +If you're using a checkpoint trained with a Diffusers training script, metadata such as the LoRA configuration, is automatically saved. When the file is loaded, the metadata is parsed to correctly configure the LoRA and avoid missing or incorrect LoRA configs. Inspect the metadata of a safetensors file by clicking on the safetensors logo logo next to the file on the Hub. -Use the [`~loaders.FromSingleFileMixin.from_single_file`] method to load a model with all the weights stored in a single safetensors file. - -```py -from diffusers import StableDiffusionPipeline - -pipeline = StableDiffusionPipeline.from_single_file( - "https://huggingface.co/WarriorMama777/OrangeMixs/blob/main/Models/AbyssOrangeMix/AbyssOrangeMix.safetensors" -) -``` - - - - -#### LoRAs - -[LoRAs](../tutorials/using_peft_for_inference) are lightweight checkpoints fine-tuned to generate images or video in a specific style. If you are using a checkpoint trained with a Diffusers training script, the LoRA configuration is automatically saved as metadata in a safetensors file. When the safetensors file is loaded, the metadata is parsed to correctly configure the LoRA and avoids missing or incorrect LoRA configurations. - -The easiest way to inspect the metadata, if available, is by clicking on the Safetensors logo next to the weights. - -
- -
- -For LoRAs that aren't trained with Diffusers, you can still save metadata with the `transformer_lora_adapter_metadata` and `text_encoder_lora_adapter_metadata` arguments in [`~loaders.FluxLoraLoaderMixin.save_lora_weights`] as long as it is a safetensors file. +Save the metadata for LoRAs that aren't trained with Diffusers with either `transformer_lora_adapter_metadata` or `unet_lora_adapter_metadata` depending on your model. For the text encoder, use the `text_encoder_lora_adapter_metadata` and `text_encoder_2_lora_adapter_metadata` arguments in [`~loaders.FluxLoraLoaderMixin.save_lora_weights`]. This is only supported for safetensors files. ```py import torch @@ -91,423 +197,88 @@ pipeline = FluxPipeline.from_pretrained( ).to("cuda") pipeline.load_lora_weights("linoyts/yarn_art_Flux_LoRA") pipeline.save_lora_weights( - transformer_lora_adapter_metadata={"r": 16, "lora_alpha": 16}, - text_encoder_lora_adapter_metadata={"r": 8, "lora_alpha": 8} + text_encoder_lora_adapter_metadata={"r": 8, "lora_alpha": 8}, + text_encoder_2_lora_adapter_metadata={"r": 8, "lora_alpha": 8} ) ``` ### ckpt -> [!WARNING] -> Pickled files may be unsafe because they can be exploited to execute malicious code. It is recommended to use safetensors files instead where possible, or convert the weights to safetensors files. +Older model weights are commonly saved with Python's [pickle](https://docs.python.org/3/library/pickle.html) utility in a ckpt file. -PyTorch's [torch.save](https://pytorch.org/docs/stable/generated/torch.save.html) function uses Python's [pickle](https://docs.python.org/3/library/pickle.html) utility to serialize and save models. These files are saved as a ckpt file and they contain the entire model's weights. +Pickled files may be unsafe because they can be exploited to execute malicious code. It is recommended to use safetensors files or convert the weights to safetensors files. -Use the [`~loaders.FromSingleFileMixin.from_single_file`] method to directly load a ckpt file. +Use [`~loaders.FromSingleFileMixin.from_single_file`] to load a ckpt file. ```py -from diffusers import StableDiffusionPipeline +from diffusers import DiffusionPipeline -pipeline = StableDiffusionPipeline.from_single_file( +pipeline = DiffusionPipeline.from_single_file( "https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5/blob/main/v1-5-pruned.ckpt" ) ``` -## Storage layout +### dduf -There are two ways model files are organized, either in a Diffusers-multifolder layout or in a single-file layout. The Diffusers-multifolder layout is the default, and each component file (text encoder, UNet, VAE) is stored in a separate subfolder. Diffusers also supports loading models from a single-file layout where all the components are bundled together. +> [!TIP] +> DDUF is an experimental file type and the API may change. Refer to the DDUF [docs](https://huggingface.co/docs/hub/dduf) to learn more. -### Diffusers-multifolder +DDUF is a file type designed to unify different diffusion model distribution methods and weight-saving formats. It is a standardized and flexible method to package all components of a diffusion model into a single file, providing a balance between the Diffusers and single-file formats. -The Diffusers-multifolder layout is the default storage layout for Diffusers. Each component's (text encoder, UNet, VAE) weights are stored in a separate subfolder. The weights can be stored as safetensors or ckpt files. - -
-
- -
multifolder layout
-
-
- -
UNet subfolder
-
-
- -To load from Diffusers-multifolder layout, use the [`~DiffusionPipeline.from_pretrained`] method. +Use the `dduf_file` argument in [`~DiffusionPipeline.from_pretrained`] to load a DDUF file. You can also load quantized dduf files as long as they are stored in the Diffusers format. ```py +import torch from diffusers import DiffusionPipeline pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", - torch_dtype=torch.float16, - variant="fp16", - use_safetensors=True, -).to("cuda") + "DDUF/FLUX.1-dev-DDUF", + dduf_file="FLUX.1-dev.dduf", + torch_dtype=torch.bfloat16, + device_map="cuda" +) ``` -Benefits of using the Diffusers-multifolder layout include: - -1. Faster to load each component file individually or in parallel. -2. Reduced memory usage because you only load the components you need. For example, models like [SDXL Turbo](https://hf.co/stabilityai/sdxl-turbo), [SDXL Lightning](https://hf.co/ByteDance/SDXL-Lightning), and [Hyper-SD](https://hf.co/ByteDance/Hyper-SD) have the same components except for the UNet. You can reuse their shared components with the [`~DiffusionPipeline.from_pipe`] method without consuming any additional memory (take a look at the [Reuse a pipeline](./loading#reuse-a-pipeline) guide) and only load the UNet. This way, you don't need to download redundant components and unnecessarily use more memory. - - ```py - import torch - from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler - - # download one model - sdxl_pipeline = StableDiffusionXLPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", - torch_dtype=torch.float16, - variant="fp16", - use_safetensors=True, - ).to("cuda") - - # switch UNet for another model - unet = UNet2DConditionModel.from_pretrained( - "stabilityai/sdxl-turbo", - subfolder="unet", - torch_dtype=torch.float16, - variant="fp16", - use_safetensors=True - ) - # reuse all the same components in new model except for the UNet - turbo_pipeline = StableDiffusionXLPipeline.from_pipe( - sdxl_pipeline, unet=unet, - ).to("cuda") - turbo_pipeline.scheduler = EulerDiscreteScheduler.from_config( - turbo_pipeline.scheduler.config, - timestep_spacing="trailing" - ) - image = turbo_pipeline( - "an astronaut riding a unicorn on mars", - num_inference_steps=1, - guidance_scale=0.0, - ).images[0] - image - ``` - -3. Reduced storage requirements because if a component, such as the SDXL [VAE](https://hf.co/madebyollin/sdxl-vae-fp16-fix), is shared across multiple models, you only need to download and store a single copy of it instead of downloading and storing it multiple times. For 10 SDXL models, this can save ~3.5GB of storage. The storage savings is even greater for newer models like PixArt Sigma, where the [text encoder](https://hf.co/PixArt-alpha/PixArt-Sigma-XL-2-1024-MS/tree/main/text_encoder) alone is ~19GB! -4. Flexibility to replace a component in the model with a newer or better version. - - ```py - from diffusers import DiffusionPipeline, AutoencoderKL - - vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16, use_safetensors=True) - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", - vae=vae, - torch_dtype=torch.float16, - variant="fp16", - use_safetensors=True, - ).to("cuda") - ``` - -5. More visibility and information about a model's components, which are stored in a [config.json](https://hf.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/unet/config.json) file in each component subfolder. - -### Single-file - -The single-file layout stores all the model weights in a single file. All the model components (text encoder, UNet, VAE) weights are kept together instead of separately in subfolders. This can be a safetensors or ckpt file. - -
- -
- -To load from a single-file layout, use the [`~loaders.FromSingleFileMixin.from_single_file`] method. +To save a pipeline as a dduf file, use the [`~huggingface_hub.export_folder_as_dduf`] utility. ```py import torch -from diffusers import StableDiffusionXLPipeline - -pipeline = StableDiffusionXLPipeline.from_single_file( - "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_base_1.0.safetensors", - torch_dtype=torch.float16, - variant="fp16", - use_safetensors=True, -).to("cuda") -``` - -Benefits of using a single-file layout include: - -1. Easy compatibility with diffusion interfaces such as [ComfyUI](https://github.com/comfyanonymous/ComfyUI) or [Automatic1111](https://github.com/AUTOMATIC1111/stable-diffusion-webui) which commonly use a single-file layout. -2. Easier to manage (download and share) a single file. - -### DDUF - -> [!WARNING] -> DDUF is an experimental file format and APIs related to it can change in the future. - -DDUF (**D**DUF **D**iffusion **U**nified **F**ormat) is a file format designed to make storing, distributing, and using diffusion models much easier. Built on the ZIP file format, DDUF offers a standardized, efficient, and flexible way to package all parts of a diffusion model into a single, easy-to-manage file. It provides a balance between Diffusers multi-folder format and the widely popular single-file format. - -Learn more details about DDUF on the Hugging Face Hub [documentation](https://huggingface.co/docs/hub/dduf). - -Pass a checkpoint to the `dduf_file` parameter to load it in [`DiffusionPipeline`]. - -```py from diffusers import DiffusionPipeline -import torch - -pipe = DiffusionPipeline.from_pretrained( - "DDUF/FLUX.1-dev-DDUF", dduf_file="FLUX.1-dev.dduf", torch_dtype=torch.bfloat16 -).to("cuda") -image = pipe( - "photo a cat holding a sign that says Diffusers", num_inference_steps=50, guidance_scale=3.5 -).images[0] -image.save("cat.png") -``` - -To save a pipeline as a `.dduf` checkpoint, use the [`~huggingface_hub.export_folder_as_dduf`] utility, which takes care of all the necessary file-level validations. - -```py from huggingface_hub import export_folder_as_dduf -from diffusers import DiffusionPipeline -import torch -pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16) +pipeline = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16) save_folder = "flux-dev" -pipe.save_pretrained("flux-dev") +pipeline.save_pretrained("flux-dev") export_folder_as_dduf("flux-dev.dduf", folder_path=save_folder) ``` -> [!TIP] -> Packaging and loading quantized checkpoints in the DDUF format is supported as long as they respect the multi-folder structure. +## Converting formats and files -## Convert layout and files +Diffusers provides scripts and methods to convert format and files to enable broader support across the diffusion ecosystem. -Diffusers provides many scripts and methods to convert storage layouts and file formats to enable broader support across the diffusion ecosystem. +Take a look at the [diffusers/scripts](https://github.com/huggingface/diffusers/tree/main/scripts) folder to find a conversion script. Scripts with `"to_diffusers` appended at the end converts a model to the Diffusers format. Each script has a specific set of arguments for configuring the conversion. Make sure you check what arguments are available. -Take a look at the [diffusers/scripts](https://github.com/huggingface/diffusers/tree/main/scripts) collection to find a script that fits your conversion needs. - -> [!TIP] -> Scripts that have "`to_diffusers`" appended at the end mean they convert a model to the Diffusers-multifolder layout. Each script has their own specific set of arguments for configuring the conversion, so make sure you check what arguments are available! - -For example, to convert a Stable Diffusion XL model stored in Diffusers-multifolder layout to a single-file layout, run the [convert_diffusers_to_original_sdxl.py](https://github.com/huggingface/diffusers/blob/main/scripts/convert_diffusers_to_original_sdxl.py) script. Provide the path to the model to convert, and the path to save the converted model to. You can optionally specify whether you want to save the model as a safetensors file and whether to save the model in half-precision. +The example below converts a model stored in Diffusers format to a single-file format. Provide the path to the model to convert and where to save the converted model. You can optionally specify what file type and data type to save the model as. ```bash python convert_diffusers_to_original_sdxl.py --model_path path/to/model/to/convert --checkpoint_path path/to/save/model/to --use_safetensors ``` -You can also save a model to Diffusers-multifolder layout with the [`~DiffusionPipeline.save_pretrained`] method. This creates a directory for you if it doesn't already exist, and it also saves the files as a safetensors file by default. +The [`~DiffusionPipeline.save_pretrained`] method also saves a model in Diffusers format and takes care of creating subfolders for each model. It saves the files as safetensor files by default. ```py -from diffusers import StableDiffusionXLPipeline +from diffusers import DiffusionPipeline -pipeline = StableDiffusionXLPipeline.from_single_file( +pipeline = DiffusionPipeline.from_single_file( "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_base_1.0.safetensors", ) pipeline.save_pretrained() ``` -Lastly, there are also Spaces, such as [SD To Diffusers](https://hf.co/spaces/diffusers/sd-to-diffusers) and [SD-XL To Diffusers](https://hf.co/spaces/diffusers/sdxl-to-diffusers), that provide a more user-friendly interface for converting models to Diffusers-multifolder layout. This is the easiest and most convenient option for converting layouts, and it'll open a PR on your model repository with the converted files. However, this option is not as reliable as running a script, and the Space may fail for more complicated models. +Finally, you can use a Space like [SD To Diffusers](https://hf.co/spaces/diffusers/sd-to-diffusers) or [SD-XL To Diffusers](https://hf.co/spaces/diffusers/sdxl-to-diffusers) to convert models to the Diffusers format. It'll open a PR on your model repository with the converted files. This is the easiest way to convert a model, but it may fail for more complicated models. Using a conversion script is more reliable. -## Single-file layout usage +## Resources -Now that you're familiar with the differences between the Diffusers-multifolder and single-file layout, this section shows you how to load models and pipeline components, customize configuration options for loading, and load local files with the [`~loaders.FromSingleFileMixin.from_single_file`] method. +- Learn more about the design decisions and why safetensor files are preferred for saving and loading model weights in the [Safetensors audited as really safe and becoming the default](https://blog.eleuther.ai/safetensors-security-audit/) blog post. -### Load a pipeline or model - -Pass the file path of the pipeline or model to the [`~loaders.FromSingleFileMixin.from_single_file`] method to load it. - - - - -```py -from diffusers import StableDiffusionXLPipeline - -ckpt_path = "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_base_1.0_0.9vae.safetensors" -pipeline = StableDiffusionXLPipeline.from_single_file(ckpt_path) -``` - - - - -```py -from diffusers import StableCascadeUNet - -ckpt_path = "https://huggingface.co/stabilityai/stable-cascade/blob/main/stage_b_lite.safetensors" -model = StableCascadeUNet.from_single_file(ckpt_path) -``` - - - - -Customize components in the pipeline by passing them directly to the [`~loaders.FromSingleFileMixin.from_single_file`] method. For example, you can use a different scheduler in a pipeline. - -```py -from diffusers import StableDiffusionXLPipeline, DDIMScheduler - -ckpt_path = "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_base_1.0_0.9vae.safetensors" -scheduler = DDIMScheduler() -pipeline = StableDiffusionXLPipeline.from_single_file(ckpt_path, scheduler=scheduler) -``` - -Or you could use a ControlNet model in the pipeline. - -```py -from diffusers import StableDiffusionControlNetPipeline, ControlNetModel - -ckpt_path = "https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5/blob/main/v1-5-pruned-emaonly.safetensors" -controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_canny") -pipeline = StableDiffusionControlNetPipeline.from_single_file(ckpt_path, controlnet=controlnet) -``` - -### Customize configuration options - -Models have a configuration file that define their attributes like the number of inputs in a UNet. Pipelines configuration options are available in the pipeline's class. For example, if you look at the [`StableDiffusionXLInstructPix2PixPipeline`] class, there is an option to scale the image latents with the `is_cosxl_edit` parameter. - -These configuration files can be found in the models Hub repository or another location from which the configuration file originated (for example, a GitHub repository or locally on your device). - - - - -> [!TIP] -> The [`~loaders.FromSingleFileMixin.from_single_file`] method automatically maps the checkpoint to the appropriate model repository, but there are cases where it is useful to use the `config` parameter. For example, if the model components in the checkpoint are different from the original checkpoint or if a checkpoint doesn't have the necessary metadata to correctly determine the configuration to use for the pipeline. - -The [`~loaders.FromSingleFileMixin.from_single_file`] method automatically determines the configuration to use from the configuration file in the model repository. You could also explicitly specify the configuration to use by providing the repository id to the `config` parameter. - -```py -from diffusers import StableDiffusionXLPipeline - -ckpt_path = "https://huggingface.co/segmind/SSD-1B/blob/main/SSD-1B.safetensors" -repo_id = "segmind/SSD-1B" - -pipeline = StableDiffusionXLPipeline.from_single_file(ckpt_path, config=repo_id) -``` - -The model loads the configuration file for the [UNet](https://huggingface.co/segmind/SSD-1B/blob/main/unet/config.json), [VAE](https://huggingface.co/segmind/SSD-1B/blob/main/vae/config.json), and [text encoder](https://huggingface.co/segmind/SSD-1B/blob/main/text_encoder/config.json) from their respective subfolders in the repository. - - - - -The [`~loaders.FromSingleFileMixin.from_single_file`] method can also load the original configuration file of a pipeline that is stored elsewhere. Pass a local path or URL of the original configuration file to the `original_config` parameter. - -```py -from diffusers import StableDiffusionXLPipeline - -ckpt_path = "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_base_1.0_0.9vae.safetensors" -original_config = "https://raw.githubusercontent.com/Stability-AI/generative-models/main/configs/inference/sd_xl_base.yaml" - -pipeline = StableDiffusionXLPipeline.from_single_file(ckpt_path, original_config=original_config) -``` - -> [!TIP] -> Diffusers attempts to infer the pipeline components based on the type signatures of the pipeline class when you use `original_config` with `local_files_only=True`, instead of fetching the configuration files from the model repository on the Hub. This prevents backward breaking changes in code that can't connect to the internet to fetch the necessary configuration files. -> -> This is not as reliable as providing a path to a local model repository with the `config` parameter, and might lead to errors during pipeline configuration. To avoid errors, run the pipeline with `local_files_only=False` once to download the appropriate pipeline configuration files to the local cache. - - - - -While the configuration files specify the pipeline or models default parameters, you can override them by providing the parameters directly to the [`~loaders.FromSingleFileMixin.from_single_file`] method. Any parameter supported by the model or pipeline class can be configured in this way. - - - - -For example, to scale the image latents in [`StableDiffusionXLInstructPix2PixPipeline`] pass the `is_cosxl_edit` parameter. - -```python -from diffusers import StableDiffusionXLInstructPix2PixPipeline - -ckpt_path = "https://huggingface.co/stabilityai/cosxl/blob/main/cosxl_edit.safetensors" -pipeline = StableDiffusionXLInstructPix2PixPipeline.from_single_file(ckpt_path, config="diffusers/sdxl-instructpix2pix-768", is_cosxl_edit=True) -``` - - - - -For example, to upcast the attention dimensions in a [`UNet2DConditionModel`] pass the `upcast_attention` parameter. - -```python -from diffusers import UNet2DConditionModel - -ckpt_path = "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_base_1.0_0.9vae.safetensors" -model = UNet2DConditionModel.from_single_file(ckpt_path, upcast_attention=True) -``` - - - - -### Local files - -In Diffusers>=v0.28.0, the [`~loaders.FromSingleFileMixin.from_single_file`] method attempts to configure a pipeline or model by inferring the model type from the keys in the checkpoint file. The inferred model type is used to determine the appropriate model repository on the Hugging Face Hub to configure the model or pipeline. - -For example, any single file checkpoint based on the Stable Diffusion XL base model will use the [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) model repository to configure the pipeline. - -But if you're working in an environment with restricted internet access, you should download the configuration files with the [`~huggingface_hub.snapshot_download`] function, and the model checkpoint with the [`~huggingface_hub.hf_hub_download`] function. By default, these files are downloaded to the Hugging Face Hub [cache directory](https://huggingface.co/docs/huggingface_hub/en/guides/manage-cache), but you can specify a preferred directory to download the files to with the `local_dir` parameter. - -Pass the configuration and checkpoint paths to the [`~loaders.FromSingleFileMixin.from_single_file`] method to load locally. - - - - -```python -from huggingface_hub import hf_hub_download, snapshot_download - -my_local_checkpoint_path = hf_hub_download( - repo_id="segmind/SSD-1B", - filename="SSD-1B.safetensors" -) - -my_local_config_path = snapshot_download( - repo_id="segmind/SSD-1B", - allow_patterns=["*.json", "**/*.json", "*.txt", "**/*.txt"] -) - -pipeline = StableDiffusionXLPipeline.from_single_file(my_local_checkpoint_path, config=my_local_config_path, local_files_only=True) -``` - - - - -```python -from huggingface_hub import hf_hub_download, snapshot_download - -my_local_checkpoint_path = hf_hub_download( - repo_id="segmind/SSD-1B", - filename="SSD-1B.safetensors" - local_dir="my_local_checkpoints" -) - -my_local_config_path = snapshot_download( - repo_id="segmind/SSD-1B", - allow_patterns=["*.json", "**/*.json", "*.txt", "**/*.txt"] - local_dir="my_local_config" -) - -pipeline = StableDiffusionXLPipeline.from_single_file(my_local_checkpoint_path, config=my_local_config_path, local_files_only=True) -``` - - - - -#### Local files without symlink - -> [!TIP] -> In huggingface_hub>=v0.23.0, the `local_dir_use_symlinks` argument isn't necessary for the [`~huggingface_hub.hf_hub_download`] and [`~huggingface_hub.snapshot_download`] functions. - -The [`~loaders.FromSingleFileMixin.from_single_file`] method relies on the [huggingface_hub](https://hf.co/docs/huggingface_hub/index) caching mechanism to fetch and store checkpoints and configuration files for models and pipelines. If you're working with a file system that does not support symlinking, you should download the checkpoint file to a local directory first, and disable symlinking with the `local_dir_use_symlink=False` parameter in the [`~huggingface_hub.hf_hub_download`] function and [`~huggingface_hub.snapshot_download`] functions. - -```python -from huggingface_hub import hf_hub_download, snapshot_download - -my_local_checkpoint_path = hf_hub_download( - repo_id="segmind/SSD-1B", - filename="SSD-1B.safetensors" - local_dir="my_local_checkpoints", - local_dir_use_symlinks=False -) -print("My local checkpoint: ", my_local_checkpoint_path) - -my_local_config_path = snapshot_download( - repo_id="segmind/SSD-1B", - allow_patterns=["*.json", "**/*.json", "*.txt", "**/*.txt"] - local_dir_use_symlinks=False, -) -print("My local config: ", my_local_config_path) -``` - -Then you can pass the local paths to the `pretrained_model_link_or_path` and `config` parameters. - -```python -pipeline = StableDiffusionXLPipeline.from_single_file(my_local_checkpoint_path, config=my_local_config_path, local_files_only=True) -``` From 76d4e416bc8952701e0b37c929bbb3253ff05f5f Mon Sep 17 00:00:00 2001 From: YiYi Xu Date: Mon, 29 Sep 2025 11:42:34 -1000 Subject: [PATCH 036/123] [modular]some small fix (#12307) * fix * add mellon node registry * style * update docstring to include more info! * support custom node mellon * HTTPErrpr -> HfHubHTTPErrpr * up * Update src/diffusers/modular_pipelines/qwenimage/node_utils.py --- .../flux/modular_pipeline.py | 2 + .../modular_pipelines/mellon_node_utils.py | 763 ++++++++++++++++++ .../modular_pipelines/modular_pipeline.py | 16 +- src/diffusers/modular_pipelines/node_utils.py | 665 --------------- .../qwenimage/before_denoise.py | 15 +- .../modular_pipelines/qwenimage/encoders.py | 2 +- .../modular_pipelines/qwenimage/inputs.py | 12 + .../qwenimage/modular_blocks.py | 64 +- .../qwenimage/modular_pipeline.py | 4 + .../modular_pipelines/qwenimage/node_utils.py | 95 +++ .../stable_diffusion_xl/before_denoise.py | 20 +- .../stable_diffusion_xl/denoise.py | 8 +- .../stable_diffusion_xl/encoders.py | 8 +- .../stable_diffusion_xl/modular_blocks.py | 59 +- .../stable_diffusion_xl/modular_pipeline.py | 2 + .../stable_diffusion_xl/node_utils.py | 99 +++ .../modular_pipelines/wan/before_denoise.py | 4 +- .../modular_pipelines/wan/denoise.py | 4 +- .../modular_pipelines/wan/encoders.py | 4 +- .../modular_pipelines/wan/modular_pipeline.py | 2 + 20 files changed, 1107 insertions(+), 741 deletions(-) create mode 100644 src/diffusers/modular_pipelines/mellon_node_utils.py delete mode 100644 src/diffusers/modular_pipelines/node_utils.py create mode 100644 src/diffusers/modular_pipelines/qwenimage/node_utils.py create mode 100644 src/diffusers/modular_pipelines/stable_diffusion_xl/node_utils.py diff --git a/src/diffusers/modular_pipelines/flux/modular_pipeline.py b/src/diffusers/modular_pipelines/flux/modular_pipeline.py index e97445d411..7d869041f2 100644 --- a/src/diffusers/modular_pipelines/flux/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/flux/modular_pipeline.py @@ -32,6 +32,8 @@ class FluxModularPipeline(ModularPipeline, FluxLoraLoaderMixin, TextualInversion """ + default_blocks_name = "FluxAutoBlocks" + @property def default_height(self): return self.default_sample_size * self.vae_scale_factor diff --git a/src/diffusers/modular_pipelines/mellon_node_utils.py b/src/diffusers/modular_pipelines/mellon_node_utils.py new file mode 100644 index 0000000000..a405aebee2 --- /dev/null +++ b/src/diffusers/modular_pipelines/mellon_node_utils.py @@ -0,0 +1,763 @@ +import json +import logging +import os + +# Simple typed wrapper for parameter overrides +from dataclasses import asdict, dataclass +from typing import Any, Dict, List, Optional, Tuple, Union + +from huggingface_hub import create_repo, hf_hub_download +from huggingface_hub.utils import ( + EntryNotFoundError, + HfHubHTTPError, + RepositoryNotFoundError, + RevisionNotFoundError, + validate_hf_hub_args, +) + +from ..utils import HUGGINGFACE_CO_RESOLVE_ENDPOINT, PushToHubMixin, extract_commit_hash +from .modular_pipeline import ModularPipelineBlocks + + +logger = logging.getLogger(__name__) + + +SUPPORTED_NODE_TYPES = {"controlnet", "vae_encoder", "denoise", "text_encoder", "decoder"} + + +# Mellon Input Parameters (runtime parameters, not models) +MELLON_INPUT_PARAMS = { + # controlnet + "control_image": { + "label": "Control Image", + "type": "image", + "display": "input", + }, + "controlnet_conditioning_scale": { + "label": "Scale", + "type": "float", + "default": 0.5, + "min": 0, + "max": 1, + }, + "control_guidance_end": { + "label": "End", + "type": "float", + "default": 1.0, + "min": 0, + "max": 1, + }, + "control_guidance_start": { + "label": "Start", + "type": "float", + "default": 0.0, + "min": 0, + "max": 1, + }, + "controlnet": { + "label": "Controlnet", + "type": "custom_controlnet", + "display": "input", + }, + "embeddings": { + "label": "Text Embeddings", + "display": "input", + "type": "embeddings", + }, + "image": { + "label": "Image", + "type": "image", + "display": "input", + }, + "negative_prompt": { + "label": "Negative Prompt", + "type": "string", + "default": "", + "display": "textarea", + }, + "prompt": { + "label": "Prompt", + "type": "string", + "default": "", + "display": "textarea", + }, + "guidance_scale": { + "label": "Guidance Scale", + "type": "float", + "display": "slider", + "default": 5, + "min": 1.0, + "max": 30.0, + "step": 0.1, + }, + "height": { + "label": "Height", + "type": "int", + "default": 1024, + "min": 64, + "step": 8, + }, + "image_latents": { + "label": "Image Latents", + "type": "latents", + "display": "input", + "onChange": {False: ["height", "width"], True: ["strength"]}, + }, + "latents": { + "label": "Latents", + "type": "latents", + "display": "input", + }, + "num_inference_steps": { + "label": "Steps", + "type": "int", + "display": "slider", + "default": 25, + "min": 1, + "max": 100, + }, + "seed": { + "label": "Seed", + "type": "int", + "display": "random", + "default": 0, + "min": 0, + "max": 4294967295, + }, + "strength": { + "label": "Strength", + "type": "float", + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + }, + "width": { + "label": "Width", + "type": "int", + "default": 1024, + "min": 64, + "step": 8, + }, + "ip_adapter": { + "label": "IP Adapter", + "type": "custom_ip_adapter", + "display": "input", + }, +} + +# Mellon Model Parameters (diffusers_auto_model types) +MELLON_MODEL_PARAMS = { + "scheduler": { + "label": "Scheduler", + "display": "input", + "type": "diffusers_auto_model", + }, + "text_encoders": { + "label": "Text Encoders", + "type": "diffusers_auto_models", + "display": "input", + }, + "unet": { + "label": "Unet", + "display": "input", + "type": "diffusers_auto_model", + "onSignal": { + "action": "signal", + "target": "guider", + }, + }, + "guider": { + "label": "Guider", + "display": "input", + "type": "custom_guider", + "onChange": {False: ["guidance_scale"], True: []}, + }, + "vae": { + "label": "VAE", + "display": "input", + "type": "diffusers_auto_model", + }, + "controlnet": { + "label": "Controlnet Model", + "type": "diffusers_auto_model", + "display": "input", + }, +} + +# Mellon Output Parameters (display = "output") +MELLON_OUTPUT_PARAMS = { + "embeddings": { + "label": "Text Embeddings", + "display": "output", + "type": "embeddings", + }, + "images": { + "label": "Images", + "type": "image", + "display": "output", + }, + "image_latents": { + "label": "Image Latents", + "type": "latents", + "display": "output", + }, + "latents": { + "label": "Latents", + "type": "latents", + "display": "output", + }, + "latents_preview": { + "label": "Latents Preview", + "display": "output", + "type": "latent", + }, + "controlnet_out": { + "label": "Controlnet", + "display": "output", + "type": "controlnet", + }, +} + + +# Default param selections per supported node_type +# from MELLON_INPUT_PARAMS / MELLON_MODEL_PARAMS / MELLON_OUTPUT_PARAMS. +NODE_TYPE_PARAMS_MAP = { + "controlnet": { + "inputs": [ + "control_image", + "controlnet_conditioning_scale", + "control_guidance_start", + "control_guidance_end", + "height", + "width", + ], + "model_inputs": [ + "controlnet", + "vae", + ], + "outputs": [ + "controlnet", + ], + "block_names": ["controlnet_vae_encoder"], + }, + "denoise": { + "inputs": [ + "embeddings", + "width", + "height", + "seed", + "num_inference_steps", + "guidance_scale", + "image_latents", + "strength", + # custom adapters coming in as inputs + "controlnet", + # ip_adapter is optional and custom; include if available + "ip_adapter", + ], + "model_inputs": [ + "unet", + "guider", + "scheduler", + ], + "outputs": [ + "latents", + "latents_preview", + ], + "block_names": ["denoise"], + }, + "vae_encoder": { + "inputs": [ + "image", + "width", + "height", + ], + "model_inputs": [ + "vae", + ], + "outputs": [ + "image_latents", + ], + "block_names": ["vae_encoder"], + }, + "text_encoder": { + "inputs": [ + "prompt", + "negative_prompt", + # optional image prompt input supported in embeddings node + "image", + ], + "model_inputs": [ + "text_encoders", + ], + "outputs": [ + "embeddings", + ], + "block_names": ["text_encoder"], + }, + "decoder": { + "inputs": [ + "latents", + ], + "model_inputs": [ + "vae", + ], + "outputs": [ + "images", + ], + "block_names": ["decode"], + }, +} + + +@dataclass(frozen=True) +class MellonParam: + name: str + label: str + type: str + display: Optional[str] = None + default: Any = None + min: Optional[float] = None + max: Optional[float] = None + step: Optional[float] = None + options: Any = None + value: Any = None + fieldOptions: Optional[Dict[str, Any]] = None + onChange: Any = None + onSignal: Any = None + _map_to_input: Any = None # the block input name this parameter maps to + + def to_dict(self) -> Dict[str, Any]: + data = asdict(self) + return {k: v for k, v in data.items() if not k.startswith("_") and v is not None} + + +@dataclass +class MellonNodeConfig(PushToHubMixin): + """ + A MellonNodeConfig is a base class to build Mellon nodes UI with modular diffusers. + + + + This is an experimental feature and is likely to change in the future. + + + """ + + inputs: List[Union[str, MellonParam]] + model_inputs: List[Union[str, MellonParam]] + outputs: List[Union[str, MellonParam]] + blocks_names: list[str] + node_type: str + config_name = "mellon_config.json" + + def __post_init__(self): + if isinstance(self.inputs, list): + self.inputs = self._resolve_params_list(self.inputs, MELLON_INPUT_PARAMS) + if isinstance(self.model_inputs, list): + self.model_inputs = self._resolve_params_list(self.model_inputs, MELLON_MODEL_PARAMS) + if isinstance(self.outputs, list): + self.outputs = self._resolve_params_list(self.outputs, MELLON_OUTPUT_PARAMS) + + @staticmethod + def _resolve_params_list( + params: List[Union[str, MellonParam]], default_map: Dict[str, Dict[str, Any]] + ) -> Dict[str, Dict[str, Any]]: + def _resolve_param( + param: Union[str, MellonParam], default_params_map: Dict[str, Dict[str, Any]] + ) -> Tuple[str, Dict[str, Any]]: + if isinstance(param, str): + if param not in default_params_map: + raise ValueError(f"Unknown param '{param}', please define a `MellonParam` object instead") + return param, default_params_map[param].copy() + elif isinstance(param, MellonParam): + param_dict = param.to_dict() + param_name = param_dict.pop("name") + return param_name, param_dict + else: + raise ValueError( + f"Unknown param type '{type(param)}', please use a string or a `MellonParam` object instead" + ) + + resolved = {} + for p in params: + logger.info(f" Resolving param: {p}") + name, cfg = _resolve_param(p, default_map) + if name in resolved: + raise ValueError(f"Duplicate param '{name}'") + resolved[name] = cfg + return resolved + + @classmethod + @validate_hf_hub_args + def load_mellon_config( + cls, + pretrained_model_name_or_path: Union[str, os.PathLike], + return_unused_kwargs=False, + return_commit_hash=False, + **kwargs, + ) -> Tuple[Dict[str, Any], Dict[str, Any]]: + r""" + Load a model or scheduler configuration. + + Parameters: + pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*): + Can be either: + + - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on + the Hub. + - A path to a *directory* (for example `./my_model_directory`) containing model weights saved with + [`~ConfigMixin.save_config`]. + + cache_dir (`Union[str, os.PathLike]`, *optional*): + Path to a directory where a downloaded pretrained model configuration is cached if the standard cache + is not used. + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force the (re-)download of the model weights and configuration files, overriding the + cached versions if they exist. + proxies (`Dict[str, str]`, *optional*): + A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', + 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. + output_loading_info(`bool`, *optional*, defaults to `False`): + Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages. + local_files_only (`bool`, *optional*, defaults to `False`): + Whether to only load local model weights and configuration files or not. If set to `True`, the model + won't be downloaded from the Hub. + token (`str` or *bool*, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from + `diffusers-cli login` (stored in `~/.huggingface`) is used. + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier + allowed by Git. + subfolder (`str`, *optional*, defaults to `""`): + The subfolder location of a model file within a larger model repository on the Hub or locally. + return_unused_kwargs (`bool`, *optional*, defaults to `False): + Whether unused keyword arguments of the config are returned. + return_commit_hash (`bool`, *optional*, defaults to `False): + Whether the `commit_hash` of the loaded configuration are returned. + + Returns: + `dict`: + A dictionary of all the parameters stored in a JSON configuration file. + + """ + cache_dir = kwargs.pop("cache_dir", None) + local_dir = kwargs.pop("local_dir", None) + local_dir_use_symlinks = kwargs.pop("local_dir_use_symlinks", "auto") + force_download = kwargs.pop("force_download", False) + proxies = kwargs.pop("proxies", None) + token = kwargs.pop("token", None) + local_files_only = kwargs.pop("local_files_only", False) + revision = kwargs.pop("revision", None) + + pretrained_model_name_or_path = str(pretrained_model_name_or_path) + + if cls.config_name is None: + raise ValueError( + "`self.config_name` is not defined. Note that one should not load a config from " + "`ConfigMixin`. Please make sure to define `config_name` in a class inheriting from `ConfigMixin`" + ) + if os.path.isfile(pretrained_model_name_or_path): + config_file = pretrained_model_name_or_path + elif os.path.isdir(pretrained_model_name_or_path): + if os.path.isfile(os.path.join(pretrained_model_name_or_path, cls.config_name)): + # Load from a PyTorch checkpoint + config_file = os.path.join(pretrained_model_name_or_path, cls.config_name) + else: + raise EnvironmentError( + f"Error no file named {cls.config_name} found in directory {pretrained_model_name_or_path}." + ) + else: + try: + # Load from URL or cache if already cached + config_file = hf_hub_download( + pretrained_model_name_or_path, + filename=cls.config_name, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + local_files_only=local_files_only, + token=token, + revision=revision, + local_dir=local_dir, + local_dir_use_symlinks=local_dir_use_symlinks, + ) + except RepositoryNotFoundError: + raise EnvironmentError( + f"{pretrained_model_name_or_path} is not a local folder and is not a valid model identifier" + " listed on 'https://huggingface.co/models'\nIf this is a private repository, make sure to pass a" + " token having permission to this repo with `token` or log in with `hf auth login`." + ) + except RevisionNotFoundError: + raise EnvironmentError( + f"{revision} is not a valid git identifier (branch name, tag name or commit id) that exists for" + " this model name. Check the model page at" + f" 'https://huggingface.co/{pretrained_model_name_or_path}' for available revisions." + ) + except EntryNotFoundError: + raise EnvironmentError( + f"{pretrained_model_name_or_path} does not appear to have a file named {cls.config_name}." + ) + except HfHubHTTPError as err: + raise EnvironmentError( + "There was a specific connection error when trying to load" + f" {pretrained_model_name_or_path}:\n{err}" + ) + except ValueError: + raise EnvironmentError( + f"We couldn't connect to '{HUGGINGFACE_CO_RESOLVE_ENDPOINT}' to load this model, couldn't find it" + f" in the cached files and it looks like {pretrained_model_name_or_path} is not the path to a" + f" directory containing a {cls.config_name} file.\nCheckout your internet connection or see how to" + " run the library in offline mode at" + " 'https://huggingface.co/docs/diffusers/installation#offline-mode'." + ) + except EnvironmentError: + raise EnvironmentError( + f"Can't load config for '{pretrained_model_name_or_path}'. If you were trying to load it from " + "'https://huggingface.co/models', make sure you don't have a local directory with the same name. " + f"Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a directory " + f"containing a {cls.config_name} file" + ) + try: + with open(config_file, "r", encoding="utf-8") as reader: + text = reader.read() + config_dict = json.loads(text) + + commit_hash = extract_commit_hash(config_file) + except (json.JSONDecodeError, UnicodeDecodeError): + raise EnvironmentError(f"It looks like the config file at '{config_file}' is not a valid JSON file.") + + if not (return_unused_kwargs or return_commit_hash): + return config_dict + + outputs = (config_dict,) + + if return_unused_kwargs: + outputs += (kwargs,) + + if return_commit_hash: + outputs += (commit_hash,) + + return outputs + + def save_mellon_config(self, save_directory: Union[str, os.PathLike], push_to_hub: bool = False, **kwargs): + """ + Save the Mellon node definition to a JSON file. + + Args: + save_directory (`str` or `os.PathLike`): + Directory where the configuration JSON file is saved (will be created if it does not exist). + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether or not to push your model to the Hugging Face Hub after saving it. You can specify the + repository you want to push to with `repo_id` (will default to the name of `save_directory` in your + namespace). + kwargs (`Dict[str, Any]`, *optional*): + Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. + """ + if os.path.isfile(save_directory): + raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file") + + os.makedirs(save_directory, exist_ok=True) + + # If we save using the predefined names, we can load using `from_config` + output_config_file = os.path.join(save_directory, self.config_name) + + self.to_json_file(output_config_file) + logger.info(f"Mellon node definition saved in {output_config_file}") + + if push_to_hub: + commit_message = kwargs.pop("commit_message", None) + private = kwargs.pop("private", None) + create_pr = kwargs.pop("create_pr", False) + token = kwargs.pop("token", None) + repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1]) + repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id + subfolder = kwargs.pop("subfolder", None) + + self._upload_folder( + save_directory, + repo_id, + token=token, + commit_message=commit_message, + create_pr=create_pr, + subfolder=subfolder, + ) + + def to_json_file(self, json_file_path: Union[str, os.PathLike]): + """ + Save the Mellon schema dictionary to a JSON file. + + Args: + json_file_path (`str` or `os.PathLike`): + Path to the JSON file to save a configuration instance's parameters. + """ + with open(json_file_path, "w", encoding="utf-8") as writer: + writer.write(self.to_json_string()) + + def to_json_string(self) -> str: + """ + Serializes this instance to a JSON string of the Mellon schema dict. + + Args: + Returns: + `str`: String containing all the attributes that make up this configuration instance in JSON format. + """ + + mellon_dict = self.to_mellon_dict() + return json.dumps(mellon_dict, indent=2, sort_keys=True) + "\n" + + def to_mellon_dict(self) -> Dict[str, Any]: + """Return a JSON-serializable dict focusing on the Mellon schema fields only. + + params is a single flat dict composed as: {**inputs, **model_inputs, **outputs}. + """ + # inputs/model_inputs/outputs are already normalized dicts + merged_params = {} + merged_params.update(self.inputs or {}) + merged_params.update(self.model_inputs or {}) + merged_params.update(self.outputs or {}) + + return { + "node_type": self.node_type, + "blocks_names": self.blocks_names, + "params": merged_params, + } + + @classmethod + def from_mellon_dict(cls, mellon_dict: Dict[str, Any]) -> "MellonNodeConfig": + """Create a config from a Mellon schema dict produced by to_mellon_dict(). + + Splits the flat params dict back into inputs/model_inputs/outputs using the known key spaces from + MELLON_INPUT_PARAMS / MELLON_MODEL_PARAMS / MELLON_OUTPUT_PARAMS. Unknown keys are treated as inputs by + default. + """ + flat_params = mellon_dict.get("params", {}) + + inputs: Dict[str, Any] = {} + model_inputs: Dict[str, Any] = {} + outputs: Dict[str, Any] = {} + + for param_name, param_dict in flat_params.items(): + if param_dict.get("display", "") == "output": + outputs[param_name] = param_dict + elif param_dict.get("type", "") in ("diffusers_auto_model", "diffusers_auto_models"): + model_inputs[param_name] = param_dict + else: + inputs[param_name] = param_dict + + return cls( + inputs=inputs, + model_inputs=model_inputs, + outputs=outputs, + blocks_names=mellon_dict.get("blocks_names", []), + node_type=mellon_dict.get("node_type"), + ) + + # YiYi Notes: not used yet + @classmethod + def from_blocks(cls, blocks: ModularPipelineBlocks, node_type: str) -> "MellonNodeConfig": + """ + Create an instance from a ModularPipeline object. If a preset exists in NODE_TYPE_PARAMS_MAP for the node_type, + use it; otherwise fall back to deriving lists from the pipeline's expected inputs/components/outputs. + """ + if node_type not in NODE_TYPE_PARAMS_MAP: + raise ValueError(f"Node type {node_type} not supported") + + blocks_names = list(blocks.sub_blocks.keys()) + + default_node_config = NODE_TYPE_PARAMS_MAP[node_type] + inputs_list: List[Union[str, MellonParam]] = default_node_config.get("inputs", []) + model_inputs_list: List[Union[str, MellonParam]] = default_node_config.get("model_inputs", []) + outputs_list: List[Union[str, MellonParam]] = default_node_config.get("outputs", []) + + for required_input_name in blocks.required_inputs: + if required_input_name not in inputs_list: + inputs_list.append( + MellonParam( + name=required_input_name, label=required_input_name, type=required_input_name, display="input" + ) + ) + + for component_spec in blocks.expected_components: + if component_spec.name not in model_inputs_list: + model_inputs_list.append( + MellonParam( + name=component_spec.name, + label=component_spec.name, + type="diffusers_auto_model", + display="input", + ) + ) + + return cls( + inputs=inputs_list, + model_inputs=model_inputs_list, + outputs=outputs_list, + blocks_names=blocks_names, + node_type=node_type, + ) + + +# Minimal modular registry for Mellon node configs +class ModularMellonNodeRegistry: + """Registry mapping (pipeline class, blocks_name) -> list of MellonNodeConfig.""" + + def __init__(self): + self._registry = {} + self._initialized = False + + def register(self, pipeline_cls: type, node_params: Dict[str, MellonNodeConfig]): + if not self._initialized: + _initialize_registry(self) + self._registry[pipeline_cls] = node_params + + def get(self, pipeline_cls: type) -> MellonNodeConfig: + if not self._initialized: + _initialize_registry(self) + return self._registry.get(pipeline_cls, None) + + def get_all(self) -> Dict[type, Dict[str, MellonNodeConfig]]: + if not self._initialized: + _initialize_registry(self) + return self._registry + + +def _register_preset_node_types( + pipeline_cls, params_map: Dict[str, Dict[str, Any]], registry: ModularMellonNodeRegistry +): + """Register all node-type presets for a given pipeline class from a params map.""" + node_configs = {} + for node_type, spec in params_map.items(): + node_config = MellonNodeConfig( + inputs=spec.get("inputs", []), + model_inputs=spec.get("model_inputs", []), + outputs=spec.get("outputs", []), + blocks_names=spec.get("block_names", []), + node_type=node_type, + ) + node_configs[node_type] = node_config + registry.register(pipeline_cls, node_configs) + + +def _initialize_registry(registry: ModularMellonNodeRegistry): + """Initialize the registry and register all available pipeline configs.""" + print("Initializing registry") + + registry._initialized = True + + try: + from .qwenimage.modular_pipeline import QwenImageModularPipeline + from .qwenimage.node_utils import QwenImage_NODE_TYPES_PARAMS_MAP + + _register_preset_node_types(QwenImageModularPipeline, QwenImage_NODE_TYPES_PARAMS_MAP, registry) + except Exception: + raise Exception("Failed to register QwenImageModularPipeline") + + try: + from .stable_diffusion_xl.modular_pipeline import StableDiffusionXLModularPipeline + from .stable_diffusion_xl.node_utils import SDXL_NODE_TYPES_PARAMS_MAP + + _register_preset_node_types(StableDiffusionXLModularPipeline, SDXL_NODE_TYPES_PARAMS_MAP, registry) + except Exception: + raise Exception("Failed to register StableDiffusionXLModularPipeline") diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 74ffc62348..206d19f173 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -51,6 +51,7 @@ if is_accelerate_available(): logger = logging.get_logger(__name__) # pylint: disable=invalid-name +# map regular pipeline to modular pipeline class name MODULAR_PIPELINE_MAPPING = OrderedDict( [ ("stable-diffusion-xl", "StableDiffusionXLModularPipeline"), @@ -61,16 +62,6 @@ MODULAR_PIPELINE_MAPPING = OrderedDict( ] ) -MODULAR_PIPELINE_BLOCKS_MAPPING = OrderedDict( - [ - ("StableDiffusionXLModularPipeline", "StableDiffusionXLAutoBlocks"), - ("WanModularPipeline", "WanAutoBlocks"), - ("FluxModularPipeline", "FluxAutoBlocks"), - ("QwenImageModularPipeline", "QwenImageAutoBlocks"), - ("QwenImageEditModularPipeline", "QwenImageEditAutoBlocks"), - ] -) - @dataclass class PipelineState: @@ -423,7 +414,7 @@ class ModularPipelineBlocks(ConfigMixin, PushToHubMixin): state.set(input_param.name, param, input_param.kwargs_type) elif input_param.kwargs_type: - # if it is a kwargs type, e.g. "guider_input_fields", it is likely to be a list of parameters + # if it is a kwargs type, e.g. "denoiser_input_fields", it is likely to be a list of parameters # we need to first find out which inputs are and loop through them. intermediate_kwargs = state.get_by_kwargs(input_param.kwargs_type) for param_name, current_value in intermediate_kwargs.items(): @@ -1454,6 +1445,7 @@ class ModularPipeline(ConfigMixin, PushToHubMixin): config_name = "modular_model_index.json" hf_device_map = None + default_blocks_name = None # YiYi TODO: add warning for passing multiple ComponentSpec/ConfigSpec with the same name def __init__( @@ -1514,7 +1506,7 @@ class ModularPipeline(ConfigMixin, PushToHubMixin): `_blocks_class_name` in the config dict """ if blocks is None: - blocks_class_name = MODULAR_PIPELINE_BLOCKS_MAPPING.get(self.__class__.__name__) + blocks_class_name = self.default_blocks_name if blocks_class_name is not None: diffusers_module = importlib.import_module("diffusers") blocks_class = getattr(diffusers_module, blocks_class_name) diff --git a/src/diffusers/modular_pipelines/node_utils.py b/src/diffusers/modular_pipelines/node_utils.py deleted file mode 100644 index 5db860c788..0000000000 --- a/src/diffusers/modular_pipelines/node_utils.py +++ /dev/null @@ -1,665 +0,0 @@ -import json -import logging -import os -from pathlib import Path -from typing import List, Optional, Tuple, Union - -import numpy as np -import PIL -import torch - -from ..configuration_utils import ConfigMixin -from ..image_processor import PipelineImageInput -from .modular_pipeline import ModularPipelineBlocks, SequentialPipelineBlocks -from .modular_pipeline_utils import InputParam - - -logger = logging.getLogger(__name__) - -# YiYi Notes: this is actually for SDXL, put it here for now -SDXL_INPUTS_SCHEMA = { - "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", - ), - "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" - ), - "eta": InputParam("eta", type_hint=float, default=0.0, description="Parameter η in the DDIM paper"), - "output_type": InputParam( - "output_type", type_hint=str, default="pil", description="Output format (pil/tensor/np.array)" - ), - "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_SCHEMA = { - "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_PARAM_SCHEMA = {**SDXL_INPUTS_SCHEMA, **SDXL_INTERMEDIATE_INPUTS_SCHEMA} - - -DEFAULT_PARAM_MAPS = { - "prompt": { - "label": "Prompt", - "type": "string", - "default": "a bear sitting in a chair drinking a milkshake", - "display": "textarea", - }, - "negative_prompt": { - "label": "Negative Prompt", - "type": "string", - "default": "deformed, ugly, wrong proportion, low res, bad anatomy, worst quality, low quality", - "display": "textarea", - }, - "num_inference_steps": { - "label": "Steps", - "type": "int", - "default": 25, - "min": 1, - "max": 1000, - }, - "seed": { - "label": "Seed", - "type": "int", - "default": 0, - "min": 0, - "display": "random", - }, - "width": { - "label": "Width", - "type": "int", - "display": "text", - "default": 1024, - "min": 8, - "max": 8192, - "step": 8, - "group": "dimensions", - }, - "height": { - "label": "Height", - "type": "int", - "display": "text", - "default": 1024, - "min": 8, - "max": 8192, - "step": 8, - "group": "dimensions", - }, - "images": { - "label": "Images", - "type": "image", - "display": "output", - }, - "image": { - "label": "Image", - "type": "image", - "display": "input", - }, -} - -DEFAULT_TYPE_MAPS = { - "int": { - "type": "int", - "default": 0, - "min": 0, - }, - "float": { - "type": "float", - "default": 0.0, - "min": 0.0, - }, - "str": { - "type": "string", - "default": "", - }, - "bool": { - "type": "boolean", - "default": False, - }, - "image": { - "type": "image", - }, -} - -DEFAULT_MODEL_KEYS = ["unet", "vae", "text_encoder", "tokenizer", "controlnet", "transformer", "image_encoder"] -DEFAULT_CATEGORY = "Modular Diffusers" -DEFAULT_EXCLUDE_MODEL_KEYS = ["processor", "feature_extractor", "safety_checker"] -DEFAULT_PARAMS_GROUPS_KEYS = { - "text_encoders": ["text_encoder", "tokenizer"], - "ip_adapter_embeds": ["ip_adapter_embeds"], - "prompt_embeddings": ["prompt_embeds"], -} - - -def get_group_name(name, group_params_keys=DEFAULT_PARAMS_GROUPS_KEYS): - """ - Get the group name for a given parameter name, if not part of a group, return None e.g. "prompt_embeds" -> - "text_embeds", "text_encoder" -> "text_encoders", "prompt" -> None - """ - if name is None: - return None - for group_name, group_keys in group_params_keys.items(): - for group_key in group_keys: - if group_key in name: - return group_name - return None - - -class ModularNode(ConfigMixin): - """ - A ModularNode is a base class to build UI nodes using diffusers. Currently only supports Mellon. It is a wrapper - around a ModularPipelineBlocks object. - - - - This is an experimental feature and is likely to change in the future. - - - """ - - config_name = "node_config.json" - - @classmethod - def from_pretrained( - cls, - pretrained_model_name_or_path: str, - trust_remote_code: Optional[bool] = None, - **kwargs, - ): - blocks = ModularPipelineBlocks.from_pretrained( - pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **kwargs - ) - return cls(blocks, **kwargs) - - def __init__(self, blocks, category=DEFAULT_CATEGORY, label=None, **kwargs): - self.blocks = blocks - - if label is None: - label = self.blocks.__class__.__name__ - # blocks param name -> mellon param name - self.name_mapping = {} - - input_params = {} - # pass or create a default param dict for each input - # e.g. for prompt, - # prompt = { - # "name": "text_input", # the name of the input in node definition, could be different from the input name in diffusers - # "label": "Prompt", - # "type": "string", - # "default": "a bear sitting in a chair drinking a milkshake", - # "display": "textarea"} - # if type is not specified, it'll be a "custom" param of its own type - # e.g. you can pass ModularNode(scheduler = {name :"scheduler"}) - # it will get this spec in node definition {"scheduler": {"label": "Scheduler", "type": "scheduler", "display": "input"}} - # name can be a dict, in that case, it is part of a "dict" input in mellon nodes, e.g. text_encoder= {name: {"text_encoders": "text_encoder"}} - inputs = self.blocks.inputs + self.blocks.intermediate_inputs - for inp in inputs: - param = kwargs.pop(inp.name, None) - if param: - # user can pass a param dict for all inputs, e.g. ModularNode(prompt = {...}) - input_params[inp.name] = param - mellon_name = param.pop("name", inp.name) - if mellon_name != inp.name: - self.name_mapping[inp.name] = mellon_name - continue - - if inp.name not in DEFAULT_PARAM_MAPS and not inp.required and not get_group_name(inp.name): - continue - - if inp.name in DEFAULT_PARAM_MAPS: - # first check if it's in the default param map, if so, directly use that - param = DEFAULT_PARAM_MAPS[inp.name].copy() - elif get_group_name(inp.name): - param = get_group_name(inp.name) - if inp.name not in self.name_mapping: - self.name_mapping[inp.name] = param - else: - # if not, check if it's in the SDXL input schema, if so, - # 1. use the type hint to determine the type - # 2. use the default param dict for the type e.g. if "steps" is a "int" type, {"steps": {"type": "int", "default": 0, "min": 0}} - if inp.type_hint is not None: - type_str = str(inp.type_hint).lower() - else: - inp_spec = SDXL_PARAM_SCHEMA.get(inp.name, None) - type_str = str(inp_spec.type_hint).lower() if inp_spec else "" - for type_key, type_param in DEFAULT_TYPE_MAPS.items(): - if type_key in type_str: - param = type_param.copy() - param["label"] = inp.name - param["display"] = "input" - break - else: - param = inp.name - # add the param dict to the inp_params dict - input_params[inp.name] = param - - component_params = {} - for comp in self.blocks.expected_components: - param = kwargs.pop(comp.name, None) - if param: - component_params[comp.name] = param - mellon_name = param.pop("name", comp.name) - if mellon_name != comp.name: - self.name_mapping[comp.name] = mellon_name - continue - - to_exclude = False - for exclude_key in DEFAULT_EXCLUDE_MODEL_KEYS: - if exclude_key in comp.name: - to_exclude = True - break - if to_exclude: - continue - - if get_group_name(comp.name): - param = get_group_name(comp.name) - if comp.name not in self.name_mapping: - self.name_mapping[comp.name] = param - elif comp.name in DEFAULT_MODEL_KEYS: - param = {"label": comp.name, "type": "diffusers_auto_model", "display": "input"} - else: - param = comp.name - # add the param dict to the model_params dict - component_params[comp.name] = param - - output_params = {} - if isinstance(self.blocks, SequentialPipelineBlocks): - last_block_name = list(self.blocks.sub_blocks.keys())[-1] - outputs = self.blocks.sub_blocks[last_block_name].intermediate_outputs - else: - outputs = self.blocks.intermediate_outputs - - for out in outputs: - param = kwargs.pop(out.name, None) - if param: - output_params[out.name] = param - mellon_name = param.pop("name", out.name) - if mellon_name != out.name: - self.name_mapping[out.name] = mellon_name - continue - - if out.name in DEFAULT_PARAM_MAPS: - param = DEFAULT_PARAM_MAPS[out.name].copy() - param["display"] = "output" - else: - group_name = get_group_name(out.name) - if group_name: - param = group_name - if out.name not in self.name_mapping: - self.name_mapping[out.name] = param - else: - param = out.name - # add the param dict to the outputs dict - output_params[out.name] = param - - if len(kwargs) > 0: - logger.warning(f"Unused kwargs: {kwargs}") - - register_dict = { - "category": category, - "label": label, - "input_params": input_params, - "component_params": component_params, - "output_params": output_params, - "name_mapping": self.name_mapping, - } - self.register_to_config(**register_dict) - - def setup(self, components_manager, collection=None): - self.pipeline = self.blocks.init_pipeline(components_manager=components_manager, collection=collection) - self._components_manager = components_manager - - @property - def mellon_config(self): - return self._convert_to_mellon_config() - - def _convert_to_mellon_config(self): - node = {} - node["label"] = self.config.label - node["category"] = self.config.category - - node_param = {} - for inp_name, inp_param in self.config.input_params.items(): - if inp_name in self.name_mapping: - mellon_name = self.name_mapping[inp_name] - else: - mellon_name = inp_name - if isinstance(inp_param, str): - param = { - "label": inp_param, - "type": inp_param, - "display": "input", - } - else: - param = inp_param - - if mellon_name not in node_param: - node_param[mellon_name] = param - else: - logger.debug(f"Input param {mellon_name} already exists in node_param, skipping {inp_name}") - - for comp_name, comp_param in self.config.component_params.items(): - if comp_name in self.name_mapping: - mellon_name = self.name_mapping[comp_name] - else: - mellon_name = comp_name - if isinstance(comp_param, str): - param = { - "label": comp_param, - "type": comp_param, - "display": "input", - } - else: - param = comp_param - - if mellon_name not in node_param: - node_param[mellon_name] = param - else: - logger.debug(f"Component param {comp_param} already exists in node_param, skipping {comp_name}") - - for out_name, out_param in self.config.output_params.items(): - if out_name in self.name_mapping: - mellon_name = self.name_mapping[out_name] - else: - mellon_name = out_name - if isinstance(out_param, str): - param = { - "label": out_param, - "type": out_param, - "display": "output", - } - else: - param = out_param - - if mellon_name not in node_param: - node_param[mellon_name] = param - else: - logger.debug(f"Output param {out_param} already exists in node_param, skipping {out_name}") - node["params"] = node_param - return node - - def save_mellon_config(self, file_path): - """ - Save the Mellon configuration to a JSON file. - - Args: - file_path (str or Path): Path where the JSON file will be saved - - Returns: - Path: Path to the saved config file - """ - file_path = Path(file_path) - - # Create directory if it doesn't exist - os.makedirs(file_path.parent, exist_ok=True) - - # Create a combined dictionary with module definition and name mapping - config = {"module": self.mellon_config, "name_mapping": self.name_mapping} - - # Save the config to file - with open(file_path, "w", encoding="utf-8") as f: - json.dump(config, f, indent=2) - - logger.info(f"Mellon config and name mapping saved to {file_path}") - - return file_path - - @classmethod - def load_mellon_config(cls, file_path): - """ - Load a Mellon configuration from a JSON file. - - Args: - file_path (str or Path): Path to the JSON file containing Mellon config - - Returns: - dict: The loaded combined configuration containing 'module' and 'name_mapping' - """ - file_path = Path(file_path) - - if not file_path.exists(): - raise FileNotFoundError(f"Config file not found: {file_path}") - - with open(file_path, "r", encoding="utf-8") as f: - config = json.load(f) - - logger.info(f"Mellon config loaded from {file_path}") - - return config - - def process_inputs(self, **kwargs): - params_components = {} - for comp_name, comp_param in self.config.component_params.items(): - logger.debug(f"component: {comp_name}") - mellon_comp_name = self.name_mapping.get(comp_name, comp_name) - if mellon_comp_name in kwargs: - if isinstance(kwargs[mellon_comp_name], dict) and comp_name in kwargs[mellon_comp_name]: - comp = kwargs[mellon_comp_name].pop(comp_name) - else: - comp = kwargs.pop(mellon_comp_name) - if comp: - params_components[comp_name] = self._components_manager.get_one(comp["model_id"]) - - params_run = {} - for inp_name, inp_param in self.config.input_params.items(): - logger.debug(f"input: {inp_name}") - mellon_inp_name = self.name_mapping.get(inp_name, inp_name) - if mellon_inp_name in kwargs: - if isinstance(kwargs[mellon_inp_name], dict) and inp_name in kwargs[mellon_inp_name]: - inp = kwargs[mellon_inp_name].pop(inp_name) - else: - inp = kwargs.pop(mellon_inp_name) - if inp is not None: - params_run[inp_name] = inp - - return_output_names = list(self.config.output_params.keys()) - - return params_components, params_run, return_output_names - - def execute(self, **kwargs): - params_components, params_run, return_output_names = self.process_inputs(**kwargs) - - self.pipeline.update_components(**params_components) - output = self.pipeline(**params_run, output=return_output_names) - return output diff --git a/src/diffusers/modular_pipelines/qwenimage/before_denoise.py b/src/diffusers/modular_pipelines/qwenimage/before_denoise.py index 738a1e5d15..606236cfe9 100644 --- a/src/diffusers/modular_pipelines/qwenimage/before_denoise.py +++ b/src/diffusers/modular_pipelines/qwenimage/before_denoise.py @@ -577,9 +577,8 @@ class QwenImageEditRoPEInputsStep(ModularPipelineBlocks): def inputs(self) -> List[InputParam]: return [ InputParam(name="batch_size", required=True), - InputParam( - name="resized_image", required=True, type_hint=torch.Tensor, description="The resized image input" - ), + InputParam(name="image_height", required=True), + InputParam(name="image_width", required=True), InputParam(name="height", required=True), InputParam(name="width", required=True), InputParam(name="prompt_embeds_mask"), @@ -612,10 +611,6 @@ class QwenImageEditRoPEInputsStep(ModularPipelineBlocks): block_state = self.get_block_state(state) # for edit, image size can be different from the target size (height/width) - image = ( - block_state.resized_image[0] if isinstance(block_state.resized_image, list) else block_state.resized_image - ) - image_width, image_height = image.size block_state.img_shapes = [ [ @@ -624,7 +619,11 @@ class QwenImageEditRoPEInputsStep(ModularPipelineBlocks): block_state.height // components.vae_scale_factor // 2, block_state.width // components.vae_scale_factor // 2, ), - (1, image_height // components.vae_scale_factor // 2, image_width // components.vae_scale_factor // 2), + ( + 1, + block_state.image_height // components.vae_scale_factor // 2, + block_state.image_width // components.vae_scale_factor // 2, + ), ] ] * block_state.batch_size diff --git a/src/diffusers/modular_pipelines/qwenimage/encoders.py b/src/diffusers/modular_pipelines/qwenimage/encoders.py index 280fa6a152..2ab83a03ee 100644 --- a/src/diffusers/modular_pipelines/qwenimage/encoders.py +++ b/src/diffusers/modular_pipelines/qwenimage/encoders.py @@ -496,7 +496,7 @@ class QwenImageEditTextEncoderStep(ModularPipelineBlocks): ) if components.requires_unconditional_embeds: - negative_prompt = block_state.negative_prompt or "" + negative_prompt = block_state.negative_prompt or " " block_state.negative_prompt_embeds, block_state.negative_prompt_embeds_mask = get_qwen_prompt_embeds_edit( components.text_encoder, components.processor, diff --git a/src/diffusers/modular_pipelines/qwenimage/inputs.py b/src/diffusers/modular_pipelines/qwenimage/inputs.py index 2b787c8238..2b229c040b 100644 --- a/src/diffusers/modular_pipelines/qwenimage/inputs.py +++ b/src/diffusers/modular_pipelines/qwenimage/inputs.py @@ -307,6 +307,13 @@ class QwenImageInputsDynamicStep(ModularPipelineBlocks): return inputs + @property + def intermediate_outputs(self) -> List[OutputParam]: + return [ + OutputParam(name="image_height", type_hint=int, description="The height of the image latents"), + OutputParam(name="image_width", type_hint=int, description="The width of the image latents"), + ] + @property def expected_components(self) -> List[ComponentSpec]: return [ @@ -327,6 +334,11 @@ class QwenImageInputsDynamicStep(ModularPipelineBlocks): block_state.height = block_state.height or height block_state.width = block_state.width or width + if not hasattr(block_state, "image_height"): + block_state.image_height = height + if not hasattr(block_state, "image_width"): + block_state.image_width = width + # 2. Patchify the image latent tensor image_latent_tensor = components.pachifier.pack_latents(image_latent_tensor) diff --git a/src/diffusers/modular_pipelines/qwenimage/modular_blocks.py b/src/diffusers/modular_pipelines/qwenimage/modular_blocks.py index a01c742fcf..9126766cc2 100644 --- a/src/diffusers/modular_pipelines/qwenimage/modular_blocks.py +++ b/src/diffusers/modular_pipelines/qwenimage/modular_blocks.py @@ -511,17 +511,42 @@ class QwenImageAutoDecodeStep(AutoPipelineBlocks): ) +class QwenImageCoreDenoiseStep(SequentialPipelineBlocks): + model_name = "qwenimage" + block_classes = [ + QwenImageAutoInputStep, + QwenImageOptionalControlNetInputStep, + QwenImageAutoBeforeDenoiseStep, + QwenImageOptionalControlNetBeforeDenoiseStep, + QwenImageAutoDenoiseStep, + ] + block_names = ["input", "controlnet_input", "before_denoise", "controlnet_before_denoise", "denoise", "decode"] + + @property + def description(self): + return ( + "Core step that performs the denoising process. \n" + + " - `QwenImageAutoInputStep` (input) standardizes the inputs for the denoising step.\n" + + " - `QwenImageOptionalControlNetInputStep` (controlnet_input) prepares the controlnet input.\n" + + " - `QwenImageAutoBeforeDenoiseStep` (before_denoise) prepares the inputs for the denoising step.\n" + + " - `QwenImageOptionalControlNetBeforeDenoiseStep` (controlnet_before_denoise) prepares the controlnet input for the denoising step.\n" + + " - `QwenImageAutoDenoiseStep` (denoise) iteratively denoises the latents.\n" + + " - `QwenImageAutoDecodeStep` (decode) decodes the latents into images.\n\n" + + "This step support text-to-image, image-to-image, inpainting, and controlnet tasks for QwenImage:\n" + + " - for image-to-image generation, you need to provide `image_latents`\n" + + " - for inpainting, you need to provide `processed_mask_image` and `image_latents`\n" + + " - to run the controlnet workflow, you need to provide `control_image_latents`\n" + + " - for text-to-image generation, all you need to provide is prompt embeddings" + ) + + ## 1.10 QwenImage/auto block & presets AUTO_BLOCKS = InsertableDict( [ ("text_encoder", QwenImageTextEncoderStep()), ("vae_encoder", QwenImageAutoVaeEncoderStep()), ("controlnet_vae_encoder", QwenImageOptionalControlNetVaeEncoderStep()), - ("input", QwenImageAutoInputStep()), - ("controlnet_input", QwenImageOptionalControlNetInputStep()), - ("before_denoise", QwenImageAutoBeforeDenoiseStep()), - ("controlnet_before_denoise", QwenImageOptionalControlNetBeforeDenoiseStep()), - ("denoise", QwenImageAutoDenoiseStep()), + ("denoise", QwenImageCoreDenoiseStep()), ("decode", QwenImageAutoDecodeStep()), ] ) @@ -699,7 +724,7 @@ class QwenImageEditAutoVaeEncoderStep(AutoPipelineBlocks): class QwenImageEditAutoInputStep(AutoPipelineBlocks): block_classes = [QwenImageInpaintInputStep, QwenImageEditInputStep] block_names = ["edit_inpaint", "edit"] - block_trigger_inputs = ["processed_mask_image", "image"] + block_trigger_inputs = ["processed_mask_image", "image_latents"] @property def description(self): @@ -800,13 +825,34 @@ class QwenImageEditAutoDenoiseStep(AutoPipelineBlocks): ## 2.7 QwenImage-Edit/auto blocks & presets + +class QwenImageEditCoreDenoiseStep(SequentialPipelineBlocks): + model_name = "qwenimage-edit" + block_classes = [ + QwenImageEditAutoInputStep, + QwenImageEditAutoBeforeDenoiseStep, + QwenImageEditAutoDenoiseStep, + ] + block_names = ["input", "before_denoise", "denoise"] + + @property + def description(self): + return ( + "Core step that performs the denoising process. \n" + + " - `QwenImageEditAutoInputStep` (input) standardizes the inputs for the denoising step.\n" + + " - `QwenImageEditAutoBeforeDenoiseStep` (before_denoise) prepares the inputs for the denoising step.\n" + + " - `QwenImageEditAutoDenoiseStep` (denoise) iteratively denoises the latents.\n\n" + + "This step support edit (img2img) and edit inpainting workflow for QwenImage Edit:\n" + + " - When `processed_mask_image` is provided, it will be used for edit inpainting task.\n" + + " - When `image_latents` is provided, it will be used for edit (img2img) task.\n" + ) + + EDIT_AUTO_BLOCKS = InsertableDict( [ ("text_encoder", QwenImageEditVLEncoderStep()), ("vae_encoder", QwenImageEditAutoVaeEncoderStep()), - ("input", QwenImageEditAutoInputStep()), - ("before_denoise", QwenImageEditAutoBeforeDenoiseStep()), - ("denoise", QwenImageEditAutoDenoiseStep()), + ("denoise", QwenImageEditCoreDenoiseStep()), ("decode", QwenImageAutoDecodeStep()), ] ) diff --git a/src/diffusers/modular_pipelines/qwenimage/modular_pipeline.py b/src/diffusers/modular_pipelines/qwenimage/modular_pipeline.py index fe9757f41b..3248d13159 100644 --- a/src/diffusers/modular_pipelines/qwenimage/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/qwenimage/modular_pipeline.py @@ -104,6 +104,8 @@ class QwenImageModularPipeline(ModularPipeline, QwenImageLoraLoaderMixin): """ + default_blocks_name = "QwenImageAutoBlocks" + @property def default_height(self): return self.default_sample_size * self.vae_scale_factor @@ -158,6 +160,8 @@ class QwenImageEditModularPipeline(ModularPipeline, QwenImageLoraLoaderMixin): """ + default_blocks_name = "QwenImageEditAutoBlocks" + # YiYi TODO: qwen edit should not provide default height/width, should be derived from the resized input image (after adjustment) produced by the resize step. @property def default_height(self): diff --git a/src/diffusers/modular_pipelines/qwenimage/node_utils.py b/src/diffusers/modular_pipelines/qwenimage/node_utils.py new file mode 100644 index 0000000000..3230ece68a --- /dev/null +++ b/src/diffusers/modular_pipelines/qwenimage/node_utils.py @@ -0,0 +1,95 @@ +# Copyright 2025 Qwen-Image Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +# mellon nodes +QwenImage_NODE_TYPES_PARAMS_MAP = { + "controlnet": { + "inputs": [ + "control_image", + "controlnet_conditioning_scale", + "control_guidance_start", + "control_guidance_end", + "height", + "width", + ], + "model_inputs": [ + "controlnet", + "vae", + ], + "outputs": [ + "controlnet_out", + ], + "block_names": ["controlnet_vae_encoder"], + }, + "denoise": { + "inputs": [ + "embeddings", + "width", + "height", + "seed", + "num_inference_steps", + "guidance_scale", + "image_latents", + "strength", + "controlnet", + ], + "model_inputs": [ + "unet", + "guider", + "scheduler", + ], + "outputs": [ + "latents", + "latents_preview", + ], + "block_names": ["denoise"], + }, + "vae_encoder": { + "inputs": [ + "image", + "width", + "height", + ], + "model_inputs": [ + "vae", + ], + "outputs": [ + "image_latents", + ], + }, + "text_encoder": { + "inputs": [ + "prompt", + "negative_prompt", + ], + "model_inputs": [ + "text_encoders", + ], + "outputs": [ + "embeddings", + ], + }, + "decoder": { + "inputs": [ + "latents", + ], + "model_inputs": [ + "vae", + ], + "outputs": [ + "images", + ], + }, +} diff --git a/src/diffusers/modular_pipelines/stable_diffusion_xl/before_denoise.py b/src/diffusers/modular_pipelines/stable_diffusion_xl/before_denoise.py index fefa622f1a..70cbf0c1c7 100644 --- a/src/diffusers/modular_pipelines/stable_diffusion_xl/before_denoise.py +++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/before_denoise.py @@ -262,37 +262,37 @@ class StableDiffusionXLInputStep(ModularPipelineBlocks): OutputParam( "prompt_embeds", type_hint=torch.Tensor, - kwargs_type="guider_input_fields", # already in intermedites state but declare here again for guider_input_fields + kwargs_type="denoiser_input_fields", # already in intermedites state but declare here again for denoiser_input_fields description="text embeddings used to guide the image generation", ), OutputParam( "negative_prompt_embeds", type_hint=torch.Tensor, - kwargs_type="guider_input_fields", # already in intermedites state but declare here again for guider_input_fields + kwargs_type="denoiser_input_fields", # already in intermedites state but declare here again for denoiser_input_fields description="negative text embeddings used to guide the image generation", ), OutputParam( "pooled_prompt_embeds", type_hint=torch.Tensor, - kwargs_type="guider_input_fields", # already in intermedites state but declare here again for guider_input_fields + kwargs_type="denoiser_input_fields", # already in intermedites state but declare here again for denoiser_input_fields description="pooled text embeddings used to guide the image generation", ), OutputParam( "negative_pooled_prompt_embeds", type_hint=torch.Tensor, - kwargs_type="guider_input_fields", # already in intermedites state but declare here again for guider_input_fields + kwargs_type="denoiser_input_fields", # already in intermedites state but declare here again for denoiser_input_fields description="negative pooled text embeddings used to guide the image generation", ), OutputParam( "ip_adapter_embeds", type_hint=List[torch.Tensor], - kwargs_type="guider_input_fields", # already in intermedites state but declare here again for guider_input_fields + kwargs_type="denoiser_input_fields", # already in intermedites state but declare here again for denoiser_input_fields description="image embeddings for IP-Adapter", ), OutputParam( "negative_ip_adapter_embeds", type_hint=List[torch.Tensor], - kwargs_type="guider_input_fields", # already in intermedites state but declare here again for guider_input_fields + kwargs_type="denoiser_input_fields", # already in intermedites state but declare here again for denoiser_input_fields description="negative image embeddings for IP-Adapter", ), ] @@ -1120,13 +1120,13 @@ class StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep(ModularPipelineB OutputParam( "add_time_ids", type_hint=torch.Tensor, - kwargs_type="guider_input_fields", + kwargs_type="denoiser_input_fields", description="The time ids to condition the denoising process", ), OutputParam( "negative_add_time_ids", type_hint=torch.Tensor, - kwargs_type="guider_input_fields", + kwargs_type="denoiser_input_fields", 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"), @@ -1331,13 +1331,13 @@ class StableDiffusionXLPrepareAdditionalConditioningStep(ModularPipelineBlocks): OutputParam( "add_time_ids", type_hint=torch.Tensor, - kwargs_type="guider_input_fields", + kwargs_type="denoiser_input_fields", description="The time ids to condition the denoising process", ), OutputParam( "negative_add_time_ids", type_hint=torch.Tensor, - kwargs_type="guider_input_fields", + kwargs_type="denoiser_input_fields", 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"), diff --git a/src/diffusers/modular_pipelines/stable_diffusion_xl/denoise.py b/src/diffusers/modular_pipelines/stable_diffusion_xl/denoise.py index a2e1420595..8a80257473 100644 --- a/src/diffusers/modular_pipelines/stable_diffusion_xl/denoise.py +++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/denoise.py @@ -183,14 +183,14 @@ class StableDiffusionXLLoopDenoiser(ModularPipelineBlocks): description="The guidance scale embedding to use for Latent Consistency Models(LCMs). Can be generated in prepare_additional_conditioning step.", ), InputParam( - kwargs_type="guider_input_fields", + kwargs_type="denoiser_input_fields", description=( "All conditional model inputs that need to be prepared with guider. " "It should contain prompt_embeds/negative_prompt_embeds, " "add_time_ids/negative_add_time_ids, " "pooled_prompt_embeds/negative_pooled_prompt_embeds, " "and ip_adapter_embeds/negative_ip_adapter_embeds (optional)." - "please add `kwargs_type=guider_input_fields` to their parameter spec (`OutputParam`) when they are created and added to the pipeline state" + "please add `kwargs_type=denoiser_input_fields` to their parameter spec (`OutputParam`) when they are created and added to the pipeline state" ), ), ] @@ -307,14 +307,14 @@ class StableDiffusionXLControlNetLoopDenoiser(ModularPipelineBlocks): description="The number of inference steps to use for the denoising process. Can be generated in set_timesteps step.", ), InputParam( - kwargs_type="guider_input_fields", + kwargs_type="denoiser_input_fields", description=( "All conditional model inputs that need to be prepared with guider. " "It should contain prompt_embeds/negative_prompt_embeds, " "add_time_ids/negative_add_time_ids, " "pooled_prompt_embeds/negative_pooled_prompt_embeds, " "and ip_adapter_embeds/negative_ip_adapter_embeds (optional)." - "please add `kwargs_type=guider_input_fields` to their parameter spec (`OutputParam`) when they are created and added to the pipeline state" + "please add `kwargs_type=denoiser_input_fields` to their parameter spec (`OutputParam`) when they are created and added to the pipeline state" ), ), InputParam( diff --git a/src/diffusers/modular_pipelines/stable_diffusion_xl/encoders.py b/src/diffusers/modular_pipelines/stable_diffusion_xl/encoders.py index 1e8921d363..90b254b6f5 100644 --- a/src/diffusers/modular_pipelines/stable_diffusion_xl/encoders.py +++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/encoders.py @@ -258,25 +258,25 @@ class StableDiffusionXLTextEncoderStep(ModularPipelineBlocks): OutputParam( "prompt_embeds", type_hint=torch.Tensor, - kwargs_type="guider_input_fields", + kwargs_type="denoiser_input_fields", description="text embeddings used to guide the image generation", ), OutputParam( "negative_prompt_embeds", type_hint=torch.Tensor, - kwargs_type="guider_input_fields", + kwargs_type="denoiser_input_fields", description="negative text embeddings used to guide the image generation", ), OutputParam( "pooled_prompt_embeds", type_hint=torch.Tensor, - kwargs_type="guider_input_fields", + kwargs_type="denoiser_input_fields", description="pooled text embeddings used to guide the image generation", ), OutputParam( "negative_pooled_prompt_embeds", type_hint=torch.Tensor, - kwargs_type="guider_input_fields", + kwargs_type="denoiser_input_fields", description="negative pooled text embeddings used to guide the image generation", ), ] diff --git a/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_blocks.py b/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_blocks.py index c9033856bc..68b5e33755 100644 --- a/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_blocks.py +++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_blocks.py @@ -82,19 +82,17 @@ class StableDiffusionXLAutoIPAdapterStep(AutoPipelineBlocks): # before_denoise: text2img class StableDiffusionXLBeforeDenoiseStep(SequentialPipelineBlocks): block_classes = [ - StableDiffusionXLInputStep, StableDiffusionXLSetTimestepsStep, StableDiffusionXLPrepareLatentsStep, StableDiffusionXLPrepareAdditionalConditioningStep, ] - block_names = ["input", "set_timesteps", "prepare_latents", "prepare_add_cond"] + block_names = ["set_timesteps", "prepare_latents", "prepare_add_cond"] @property def description(self): return ( "Before denoise step that prepare the inputs for the denoise step.\n" + "This is a sequential pipeline blocks:\n" - + " - `StableDiffusionXLInputStep` is used to adjust the batch size of the model inputs\n" + " - `StableDiffusionXLSetTimestepsStep` is used to set the timesteps\n" + " - `StableDiffusionXLPrepareLatentsStep` is used to prepare the latents\n" + " - `StableDiffusionXLPrepareAdditionalConditioningStep` is used to prepare the additional conditioning\n" @@ -104,19 +102,17 @@ class StableDiffusionXLBeforeDenoiseStep(SequentialPipelineBlocks): # before_denoise: img2img class StableDiffusionXLImg2ImgBeforeDenoiseStep(SequentialPipelineBlocks): block_classes = [ - StableDiffusionXLInputStep, StableDiffusionXLImg2ImgSetTimestepsStep, StableDiffusionXLImg2ImgPrepareLatentsStep, StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep, ] - block_names = ["input", "set_timesteps", "prepare_latents", "prepare_add_cond"] + block_names = ["set_timesteps", "prepare_latents", "prepare_add_cond"] @property def description(self): return ( "Before denoise step that prepare the inputs for the denoise step for img2img task.\n" + "This is a sequential pipeline blocks:\n" - + " - `StableDiffusionXLInputStep` is used to adjust the batch size of the model inputs\n" + " - `StableDiffusionXLImg2ImgSetTimestepsStep` is used to set the timesteps\n" + " - `StableDiffusionXLImg2ImgPrepareLatentsStep` is used to prepare the latents\n" + " - `StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep` is used to prepare the additional conditioning\n" @@ -126,19 +122,17 @@ class StableDiffusionXLImg2ImgBeforeDenoiseStep(SequentialPipelineBlocks): # before_denoise: inpainting class StableDiffusionXLInpaintBeforeDenoiseStep(SequentialPipelineBlocks): block_classes = [ - StableDiffusionXLInputStep, StableDiffusionXLImg2ImgSetTimestepsStep, StableDiffusionXLInpaintPrepareLatentsStep, StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep, ] - block_names = ["input", "set_timesteps", "prepare_latents", "prepare_add_cond"] + block_names = ["set_timesteps", "prepare_latents", "prepare_add_cond"] @property def description(self): return ( "Before denoise step that prepare the inputs for the denoise step for inpainting task.\n" + "This is a sequential pipeline blocks:\n" - + " - `StableDiffusionXLInputStep` is used to adjust the batch size of the model inputs\n" + " - `StableDiffusionXLImg2ImgSetTimestepsStep` is used to set the timesteps\n" + " - `StableDiffusionXLInpaintPrepareLatentsStep` is used to prepare the latents\n" + " - `StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep` is used to prepare the additional conditioning\n" @@ -255,25 +249,48 @@ class StableDiffusionXLAutoDecodeStep(AutoPipelineBlocks): ) +class StableDiffusionXLCoreDenoiseStep(SequentialPipelineBlocks): + block_classes = [ + StableDiffusionXLInputStep, + StableDiffusionXLAutoBeforeDenoiseStep, + StableDiffusionXLAutoControlNetInputStep, + StableDiffusionXLAutoDenoiseStep, + ] + block_names = ["input", "before_denoise", "controlnet_input", "denoise"] + + @property + def description(self): + return ( + "Core step that performs the denoising process. \n" + + " - `StableDiffusionXLInputStep` (input) standardizes the inputs for the denoising step.\n" + + " - `StableDiffusionXLAutoBeforeDenoiseStep` (before_denoise) prepares the inputs for the denoising step.\n" + + " - `StableDiffusionXLAutoControlNetInputStep` (controlnet_input) prepares the controlnet input.\n" + + " - `StableDiffusionXLAutoDenoiseStep` (denoise) iteratively denoises the latents.\n\n" + + "This step support text-to-image, image-to-image, inpainting, with or without controlnet/controlnet_union/ip_adapter for Stable Diffusion XL:\n" + + "- for image-to-image generation, you need to provide `image_latents`\n" + + "- for inpainting, you need to provide `mask_image` and `image_latents`\n" + + "- to run the controlnet workflow, you need to provide `control_image`\n" + + "- to run the controlnet_union workflow, you need to provide `control_image` and `control_mode`\n" + + "- to run the ip_adapter workflow, you need to load ip_adapter into your unet and provide `ip_adapter_embeds`\n" + + "- for text-to-image generation, all you need to provide is prompt embeddings\n" + ) + + # ip-adapter, controlnet, text2img, img2img, inpainting class StableDiffusionXLAutoBlocks(SequentialPipelineBlocks): block_classes = [ StableDiffusionXLTextEncoderStep, StableDiffusionXLAutoIPAdapterStep, StableDiffusionXLAutoVaeEncoderStep, - StableDiffusionXLAutoBeforeDenoiseStep, - StableDiffusionXLAutoControlNetInputStep, - StableDiffusionXLAutoDenoiseStep, + StableDiffusionXLCoreDenoiseStep, StableDiffusionXLAutoDecodeStep, ] block_names = [ "text_encoder", "ip_adapter", - "image_encoder", - "before_denoise", - "controlnet_input", + "vae_encoder", "denoise", - "decoder", + "decode", ] @property @@ -321,7 +338,7 @@ TEXT2IMAGE_BLOCKS = InsertableDict( IMAGE2IMAGE_BLOCKS = InsertableDict( [ ("text_encoder", StableDiffusionXLTextEncoderStep), - ("image_encoder", StableDiffusionXLVaeEncoderStep), + ("vae_encoder", StableDiffusionXLVaeEncoderStep), ("input", StableDiffusionXLInputStep), ("set_timesteps", StableDiffusionXLImg2ImgSetTimestepsStep), ("prepare_latents", StableDiffusionXLImg2ImgPrepareLatentsStep), @@ -334,7 +351,7 @@ IMAGE2IMAGE_BLOCKS = InsertableDict( INPAINT_BLOCKS = InsertableDict( [ ("text_encoder", StableDiffusionXLTextEncoderStep), - ("image_encoder", StableDiffusionXLInpaintVaeEncoderStep), + ("vae_encoder", StableDiffusionXLInpaintVaeEncoderStep), ("input", StableDiffusionXLInputStep), ("set_timesteps", StableDiffusionXLImg2ImgSetTimestepsStep), ("prepare_latents", StableDiffusionXLInpaintPrepareLatentsStep), @@ -361,10 +378,8 @@ AUTO_BLOCKS = InsertableDict( [ ("text_encoder", StableDiffusionXLTextEncoderStep), ("ip_adapter", StableDiffusionXLAutoIPAdapterStep), - ("image_encoder", StableDiffusionXLAutoVaeEncoderStep), - ("before_denoise", StableDiffusionXLAutoBeforeDenoiseStep), - ("controlnet_input", StableDiffusionXLAutoControlNetInputStep), - ("denoise", StableDiffusionXLAutoDenoiseStep), + ("vae_encoder", StableDiffusionXLAutoVaeEncoderStep), + ("denoise", StableDiffusionXLCoreDenoiseStep), ("decode", StableDiffusionXLAutoDecodeStep), ] ) diff --git a/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_pipeline.py b/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_pipeline.py index e84f5cad1a..29a717f72e 100644 --- a/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_pipeline.py @@ -54,6 +54,8 @@ class StableDiffusionXLModularPipeline( """ + default_blocks_name = "StableDiffusionXLAutoBlocks" + @property def default_height(self): return self.default_sample_size * self.vae_scale_factor diff --git a/src/diffusers/modular_pipelines/stable_diffusion_xl/node_utils.py b/src/diffusers/modular_pipelines/stable_diffusion_xl/node_utils.py new file mode 100644 index 0000000000..3e788bf947 --- /dev/null +++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/node_utils.py @@ -0,0 +1,99 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +SDXL_NODE_TYPES_PARAMS_MAP = { + "controlnet": { + "inputs": [ + "control_image", + "controlnet_conditioning_scale", + "control_guidance_start", + "control_guidance_end", + "height", + "width", + ], + "model_inputs": [ + "controlnet", + ], + "outputs": [ + "controlnet_out", + ], + "block_names": [None], + }, + "denoise": { + "inputs": [ + "embeddings", + "width", + "height", + "seed", + "num_inference_steps", + "guidance_scale", + "image_latents", + "strength", + # custom adapters coming in as inputs + "controlnet", + # ip_adapter is optional and custom; include if available + "ip_adapter", + ], + "model_inputs": [ + "unet", + "guider", + "scheduler", + ], + "outputs": [ + "latents", + "latents_preview", + ], + "block_names": ["denoise"], + }, + "vae_encoder": { + "inputs": [ + "image", + "width", + "height", + ], + "model_inputs": [ + "vae", + ], + "outputs": [ + "image_latents", + ], + "block_names": ["vae_encoder"], + }, + "text_encoder": { + "inputs": [ + "prompt", + "negative_prompt", + ], + "model_inputs": [ + "text_encoders", + ], + "outputs": [ + "embeddings", + ], + "block_names": ["text_encoder"], + }, + "decoder": { + "inputs": [ + "latents", + ], + "model_inputs": [ + "vae", + ], + "outputs": [ + "images", + ], + "block_names": ["decode"], + }, +} diff --git a/src/diffusers/modular_pipelines/wan/before_denoise.py b/src/diffusers/modular_pipelines/wan/before_denoise.py index 2b9889f877..d48f678edd 100644 --- a/src/diffusers/modular_pipelines/wan/before_denoise.py +++ b/src/diffusers/modular_pipelines/wan/before_denoise.py @@ -146,13 +146,13 @@ class WanInputStep(ModularPipelineBlocks): OutputParam( "prompt_embeds", type_hint=torch.Tensor, - kwargs_type="guider_input_fields", # already in intermedites state but declare here again for guider_input_fields + kwargs_type="denoiser_input_fields", # already in intermedites state but declare here again for denoiser_input_fields description="text embeddings used to guide the image generation", ), OutputParam( "negative_prompt_embeds", type_hint=torch.Tensor, - kwargs_type="guider_input_fields", # already in intermedites state but declare here again for guider_input_fields + kwargs_type="denoiser_input_fields", # already in intermedites state but declare here again for denoiser_input_fields description="negative text embeddings used to guide the image generation", ), ] diff --git a/src/diffusers/modular_pipelines/wan/denoise.py b/src/diffusers/modular_pipelines/wan/denoise.py index 5f578609c2..66c51493bd 100644 --- a/src/diffusers/modular_pipelines/wan/denoise.py +++ b/src/diffusers/modular_pipelines/wan/denoise.py @@ -79,11 +79,11 @@ class WanLoopDenoiser(ModularPipelineBlocks): description="The number of inference steps to use for the denoising process. Can be generated in set_timesteps step.", ), InputParam( - kwargs_type="guider_input_fields", + kwargs_type="denoiser_input_fields", description=( "All conditional model inputs that need to be prepared with guider. " "It should contain prompt_embeds/negative_prompt_embeds. " - "Please add `kwargs_type=guider_input_fields` to their parameter spec (`OutputParam`) when they are created and added to the pipeline state" + "Please add `kwargs_type=denoiser_input_fields` to their parameter spec (`OutputParam`) when they are created and added to the pipeline state" ), ), ] diff --git a/src/diffusers/modular_pipelines/wan/encoders.py b/src/diffusers/modular_pipelines/wan/encoders.py index a0bf76b99b..cb2fc24238 100644 --- a/src/diffusers/modular_pipelines/wan/encoders.py +++ b/src/diffusers/modular_pipelines/wan/encoders.py @@ -89,13 +89,13 @@ class WanTextEncoderStep(ModularPipelineBlocks): OutputParam( "prompt_embeds", type_hint=torch.Tensor, - kwargs_type="guider_input_fields", + kwargs_type="denoiser_input_fields", description="text embeddings used to guide the image generation", ), OutputParam( "negative_prompt_embeds", type_hint=torch.Tensor, - kwargs_type="guider_input_fields", + kwargs_type="denoiser_input_fields", description="negative text embeddings used to guide the image generation", ), ] diff --git a/src/diffusers/modular_pipelines/wan/modular_pipeline.py b/src/diffusers/modular_pipelines/wan/modular_pipeline.py index 4d86e0d08e..da4aada438 100644 --- a/src/diffusers/modular_pipelines/wan/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/wan/modular_pipeline.py @@ -37,6 +37,8 @@ class WanModularPipeline( """ + default_blocks_name = "WanAutoBlocks" + @property def default_height(self): return self.default_sample_height * self.vae_scale_factor_spatial From 20fd00b14b6ad0ec3fa1b67e26323e135bf9fe8a Mon Sep 17 00:00:00 2001 From: Dhruv Nair Date: Tue, 30 Sep 2025 09:58:34 +0200 Subject: [PATCH 037/123] [Tests] Add single file tester mixin for Models and remove unittest dependency (#12352) * update * update * update * update * update --- .../single_file/single_file_testing_utils.py | 91 +++++++++++++++++++ tests/single_file/test_lumina2_transformer.py | 41 +-------- .../test_model_autoencoder_dc_single_file.py | 32 +------ .../test_model_controlnet_single_file.py | 33 +------ ...test_model_flux_transformer_single_file.py | 38 +------- .../test_model_motion_adapter_single_file.py | 3 +- .../test_model_sd_cascade_unet_single_file.py | 9 +- .../single_file/test_model_vae_single_file.py | 32 +------ .../test_model_wan_autoencoder_single_file.py | 32 +------ ...est_model_wan_transformer3d_single_file.py | 58 +----------- tests/single_file/test_sana_transformer.py | 42 +-------- ...iffusion_controlnet_img2img_single_file.py | 9 +- ...iffusion_controlnet_inpaint_single_file.py | 14 ++- ...stable_diffusion_controlnet_single_file.py | 9 +- ...st_stable_diffusion_img2img_single_file.py | 17 ++-- ...st_stable_diffusion_inpaint_single_file.py | 22 ++--- .../test_stable_diffusion_single_file.py | 25 ++--- ...st_stable_diffusion_upscale_single_file.py | 9 +- ...stable_diffusion_xl_adapter_single_file.py | 9 +- ...ble_diffusion_xl_controlnet_single_file.py | 9 +- ...stable_diffusion_xl_img2img_single_file.py | 11 +-- ...st_stable_diffusion_xl_instruct_pix2pix.py | 9 +- .../test_stable_diffusion_xl_single_file.py | 9 +- 23 files changed, 173 insertions(+), 390 deletions(-) diff --git a/tests/single_file/single_file_testing_utils.py b/tests/single_file/single_file_testing_utils.py index 3510d3371c..52fd2f5bfc 100644 --- a/tests/single_file/single_file_testing_utils.py +++ b/tests/single_file/single_file_testing_utils.py @@ -1,3 +1,4 @@ +import gc import tempfile from io import BytesIO @@ -9,7 +10,10 @@ from diffusers.loaders.single_file_utils import _extract_repo_id_and_weights_nam from diffusers.models.attention_processor import AttnProcessor from ..testing_utils import ( + backend_empty_cache, + nightly, numpy_cosine_similarity_distance, + require_torch_accelerator, torch_device, ) @@ -47,6 +51,93 @@ def download_diffusers_config(repo_id, tmpdir): return path +@nightly +@require_torch_accelerator +class SingleFileModelTesterMixin: + def setup_method(self): + gc.collect() + backend_empty_cache(torch_device) + + def teardown_method(self): + gc.collect() + backend_empty_cache(torch_device) + + def test_single_file_model_config(self): + pretrained_kwargs = {} + single_file_kwargs = {} + + if hasattr(self, "subfolder") and self.subfolder: + pretrained_kwargs["subfolder"] = self.subfolder + + if hasattr(self, "torch_dtype") and self.torch_dtype: + pretrained_kwargs["torch_dtype"] = self.torch_dtype + single_file_kwargs["torch_dtype"] = self.torch_dtype + + model = self.model_class.from_pretrained(self.repo_id, **pretrained_kwargs) + model_single_file = self.model_class.from_single_file(self.ckpt_path, **single_file_kwargs) + + PARAMS_TO_IGNORE = ["torch_dtype", "_name_or_path", "_use_default_values", "_diffusers_version"] + for param_name, param_value in model_single_file.config.items(): + if param_name in PARAMS_TO_IGNORE: + continue + assert model.config[param_name] == param_value, ( + f"{param_name} differs between pretrained loading and single file loading" + ) + + def test_single_file_model_parameters(self): + pretrained_kwargs = {} + single_file_kwargs = {} + + if hasattr(self, "subfolder") and self.subfolder: + pretrained_kwargs["subfolder"] = self.subfolder + + if hasattr(self, "torch_dtype") and self.torch_dtype: + pretrained_kwargs["torch_dtype"] = self.torch_dtype + single_file_kwargs["torch_dtype"] = self.torch_dtype + + model = self.model_class.from_pretrained(self.repo_id, **pretrained_kwargs) + model_single_file = self.model_class.from_single_file(self.ckpt_path, **single_file_kwargs) + + state_dict = model.state_dict() + state_dict_single_file = model_single_file.state_dict() + + assert set(state_dict.keys()) == set(state_dict_single_file.keys()), ( + "Model parameters keys differ between pretrained and single file loading" + ) + + for key in state_dict.keys(): + param = state_dict[key] + param_single_file = state_dict_single_file[key] + + assert param.shape == param_single_file.shape, ( + f"Parameter shape mismatch for {key}: " + f"pretrained {param.shape} vs single file {param_single_file.shape}" + ) + + assert torch.allclose(param, param_single_file, rtol=1e-5, atol=1e-5), ( + f"Parameter values differ for {key}: " + f"max difference {torch.max(torch.abs(param - param_single_file)).item()}" + ) + + def test_checkpoint_altered_keys_loading(self): + # Test loading with checkpoints that have altered keys + if not hasattr(self, "alternate_keys_ckpt_paths") or not self.alternate_keys_ckpt_paths: + return + + for ckpt_path in self.alternate_keys_ckpt_paths: + backend_empty_cache(torch_device) + + single_file_kwargs = {} + if hasattr(self, "torch_dtype") and self.torch_dtype: + single_file_kwargs["torch_dtype"] = self.torch_dtype + + model = self.model_class.from_single_file(ckpt_path, **single_file_kwargs) + + del model + gc.collect() + backend_empty_cache(torch_device) + + class SDSingleFileTesterMixin: single_file_kwargs = {} diff --git a/tests/single_file/test_lumina2_transformer.py b/tests/single_file/test_lumina2_transformer.py index 99d9b71395..bb5a0bf473 100644 --- a/tests/single_file/test_lumina2_transformer.py +++ b/tests/single_file/test_lumina2_transformer.py @@ -13,26 +13,21 @@ # See the License for the specific language governing permissions and # limitations under the License. -import gc -import unittest from diffusers import ( Lumina2Transformer2DModel, ) from ..testing_utils import ( - backend_empty_cache, enable_full_determinism, - require_torch_accelerator, - torch_device, ) +from .single_file_testing_utils import SingleFileModelTesterMixin enable_full_determinism() -@require_torch_accelerator -class Lumina2Transformer2DModelSingleFileTests(unittest.TestCase): +class TestLumina2Transformer2DModelSingleFile(SingleFileModelTesterMixin): model_class = Lumina2Transformer2DModel ckpt_path = "https://huggingface.co/Comfy-Org/Lumina_Image_2.0_Repackaged/blob/main/split_files/diffusion_models/lumina_2_model_bf16.safetensors" alternate_keys_ckpt_paths = [ @@ -40,34 +35,4 @@ class Lumina2Transformer2DModelSingleFileTests(unittest.TestCase): ] repo_id = "Alpha-VLLM/Lumina-Image-2.0" - - def setUp(self): - super().setUp() - gc.collect() - backend_empty_cache(torch_device) - - def tearDown(self): - super().tearDown() - gc.collect() - backend_empty_cache(torch_device) - - def test_single_file_components(self): - model = self.model_class.from_pretrained(self.repo_id, subfolder="transformer") - model_single_file = self.model_class.from_single_file(self.ckpt_path) - - PARAMS_TO_IGNORE = ["torch_dtype", "_name_or_path", "_use_default_values", "_diffusers_version"] - for param_name, param_value in model_single_file.config.items(): - if param_name in PARAMS_TO_IGNORE: - continue - assert model.config[param_name] == param_value, ( - f"{param_name} differs between single file loading and pretrained loading" - ) - - def test_checkpoint_loading(self): - for ckpt_path in self.alternate_keys_ckpt_paths: - backend_empty_cache(torch_device) - model = self.model_class.from_single_file(ckpt_path) - - del model - gc.collect() - backend_empty_cache(torch_device) + subfolder = "transformer" diff --git a/tests/single_file/test_model_autoencoder_dc_single_file.py b/tests/single_file/test_model_autoencoder_dc_single_file.py index 5195f8e52f..444ca40469 100644 --- a/tests/single_file/test_model_autoencoder_dc_single_file.py +++ b/tests/single_file/test_model_autoencoder_dc_single_file.py @@ -13,8 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import gc -import unittest import torch @@ -23,38 +21,24 @@ from diffusers import ( ) from ..testing_utils import ( - backend_empty_cache, enable_full_determinism, load_hf_numpy, numpy_cosine_similarity_distance, - require_torch_accelerator, - slow, torch_device, ) +from .single_file_testing_utils import SingleFileModelTesterMixin enable_full_determinism() -@slow -@require_torch_accelerator -class AutoencoderDCSingleFileTests(unittest.TestCase): +class TestAutoencoderDCSingleFile(SingleFileModelTesterMixin): model_class = AutoencoderDC ckpt_path = "https://huggingface.co/mit-han-lab/dc-ae-f32c32-sana-1.0/blob/main/model.safetensors" repo_id = "mit-han-lab/dc-ae-f32c32-sana-1.0-diffusers" main_input_name = "sample" base_precision = 1e-2 - def setUp(self): - super().setUp() - gc.collect() - backend_empty_cache(torch_device) - - def tearDown(self): - super().tearDown() - gc.collect() - backend_empty_cache(torch_device) - def get_file_format(self, seed, shape): return f"gaussian_noise_s={seed}_shape={'_'.join([str(s) for s in shape])}.npy" @@ -80,18 +64,6 @@ class AutoencoderDCSingleFileTests(unittest.TestCase): assert numpy_cosine_similarity_distance(output_slice_1, output_slice_2) < 1e-4 - def test_single_file_components(self): - model = self.model_class.from_pretrained(self.repo_id) - model_single_file = self.model_class.from_single_file(self.ckpt_path) - - PARAMS_TO_IGNORE = ["torch_dtype", "_name_or_path", "_use_default_values", "_diffusers_version"] - for param_name, param_value in model_single_file.config.items(): - if param_name in PARAMS_TO_IGNORE: - continue - assert model.config[param_name] == param_value, ( - f"{param_name} differs between pretrained loading and single file loading" - ) - def test_single_file_in_type_variant_components(self): # `in` variant checkpoints require passing in a `config` parameter # in order to set the scaling factor correctly. diff --git a/tests/single_file/test_model_controlnet_single_file.py b/tests/single_file/test_model_controlnet_single_file.py index e5214fe3f2..2fa81fe3ae 100644 --- a/tests/single_file/test_model_controlnet_single_file.py +++ b/tests/single_file/test_model_controlnet_single_file.py @@ -13,8 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import gc -import unittest import torch @@ -23,46 +21,19 @@ from diffusers import ( ) from ..testing_utils import ( - backend_empty_cache, enable_full_determinism, - require_torch_accelerator, - slow, - torch_device, ) +from .single_file_testing_utils import SingleFileModelTesterMixin enable_full_determinism() -@slow -@require_torch_accelerator -class ControlNetModelSingleFileTests(unittest.TestCase): +class TestControlNetModelSingleFile(SingleFileModelTesterMixin): model_class = ControlNetModel ckpt_path = "https://huggingface.co/lllyasviel/ControlNet-v1-1/blob/main/control_v11p_sd15_canny.pth" repo_id = "lllyasviel/control_v11p_sd15_canny" - def setUp(self): - super().setUp() - gc.collect() - backend_empty_cache(torch_device) - - def tearDown(self): - super().tearDown() - gc.collect() - backend_empty_cache(torch_device) - - def test_single_file_components(self): - model = self.model_class.from_pretrained(self.repo_id) - model_single_file = self.model_class.from_single_file(self.ckpt_path) - - PARAMS_TO_IGNORE = ["torch_dtype", "_name_or_path", "_use_default_values", "_diffusers_version"] - for param_name, param_value in model_single_file.config.items(): - if param_name in PARAMS_TO_IGNORE: - continue - assert model.config[param_name] == param_value, ( - f"{param_name} differs between single file loading and pretrained loading" - ) - def test_single_file_arguments(self): model_default = self.model_class.from_single_file(self.ckpt_path) diff --git a/tests/single_file/test_model_flux_transformer_single_file.py b/tests/single_file/test_model_flux_transformer_single_file.py index 8290c339b9..0642a71c57 100644 --- a/tests/single_file/test_model_flux_transformer_single_file.py +++ b/tests/single_file/test_model_flux_transformer_single_file.py @@ -14,7 +14,6 @@ # limitations under the License. import gc -import unittest from diffusers import ( FluxTransformer2DModel, @@ -23,52 +22,21 @@ from diffusers import ( from ..testing_utils import ( backend_empty_cache, enable_full_determinism, - require_torch_accelerator, torch_device, ) +from .single_file_testing_utils import SingleFileModelTesterMixin enable_full_determinism() -@require_torch_accelerator -class FluxTransformer2DModelSingleFileTests(unittest.TestCase): +class TestFluxTransformer2DModelSingleFile(SingleFileModelTesterMixin): model_class = FluxTransformer2DModel ckpt_path = "https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/flux1-dev.safetensors" alternate_keys_ckpt_paths = ["https://huggingface.co/Comfy-Org/flux1-dev/blob/main/flux1-dev-fp8.safetensors"] repo_id = "black-forest-labs/FLUX.1-dev" - - def setUp(self): - super().setUp() - gc.collect() - backend_empty_cache(torch_device) - - def tearDown(self): - super().tearDown() - gc.collect() - backend_empty_cache(torch_device) - - def test_single_file_components(self): - model = self.model_class.from_pretrained(self.repo_id, subfolder="transformer") - model_single_file = self.model_class.from_single_file(self.ckpt_path) - - PARAMS_TO_IGNORE = ["torch_dtype", "_name_or_path", "_use_default_values", "_diffusers_version"] - for param_name, param_value in model_single_file.config.items(): - if param_name in PARAMS_TO_IGNORE: - continue - assert model.config[param_name] == param_value, ( - f"{param_name} differs between single file loading and pretrained loading" - ) - - def test_checkpoint_loading(self): - for ckpt_path in self.alternate_keys_ckpt_paths: - backend_empty_cache(torch_device) - model = self.model_class.from_single_file(ckpt_path) - - del model - gc.collect() - backend_empty_cache(torch_device) + subfolder = "transformer" def test_device_map_cuda(self): backend_empty_cache(torch_device) diff --git a/tests/single_file/test_model_motion_adapter_single_file.py b/tests/single_file/test_model_motion_adapter_single_file.py index 7aaf4b577e..a047c81b47 100644 --- a/tests/single_file/test_model_motion_adapter_single_file.py +++ b/tests/single_file/test_model_motion_adapter_single_file.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import unittest from diffusers import ( MotionAdapter, @@ -27,7 +26,7 @@ from ..testing_utils import ( enable_full_determinism() -class MotionAdapterSingleFileTests(unittest.TestCase): +class MotionAdapterSingleFileTests: model_class = MotionAdapter def test_single_file_components_version_v1_5(self): diff --git a/tests/single_file/test_model_sd_cascade_unet_single_file.py b/tests/single_file/test_model_sd_cascade_unet_single_file.py index a5ec9dba30..7472122710 100644 --- a/tests/single_file/test_model_sd_cascade_unet_single_file.py +++ b/tests/single_file/test_model_sd_cascade_unet_single_file.py @@ -14,7 +14,6 @@ # limitations under the License. import gc -import unittest import torch @@ -37,14 +36,12 @@ enable_full_determinism() @slow @require_torch_accelerator -class StableCascadeUNetSingleFileTest(unittest.TestCase): - def setUp(self): - super().setUp() +class StableCascadeUNetSingleFileTest: + def setup_method(self): gc.collect() backend_empty_cache(torch_device) - def tearDown(self): - super().tearDown() + def teardown_method(self): gc.collect() backend_empty_cache(torch_device) diff --git a/tests/single_file/test_model_vae_single_file.py b/tests/single_file/test_model_vae_single_file.py index 3b9e619f13..9198d9b163 100644 --- a/tests/single_file/test_model_vae_single_file.py +++ b/tests/single_file/test_model_vae_single_file.py @@ -13,8 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import gc -import unittest import torch @@ -23,22 +21,18 @@ from diffusers import ( ) from ..testing_utils import ( - backend_empty_cache, enable_full_determinism, load_hf_numpy, numpy_cosine_similarity_distance, - require_torch_accelerator, - slow, torch_device, ) +from .single_file_testing_utils import SingleFileModelTesterMixin enable_full_determinism() -@slow -@require_torch_accelerator -class AutoencoderKLSingleFileTests(unittest.TestCase): +class TestAutoencoderKLSingleFile(SingleFileModelTesterMixin): model_class = AutoencoderKL ckpt_path = ( "https://huggingface.co/stabilityai/sd-vae-ft-mse-original/blob/main/vae-ft-mse-840000-ema-pruned.safetensors" @@ -47,16 +41,6 @@ class AutoencoderKLSingleFileTests(unittest.TestCase): main_input_name = "sample" base_precision = 1e-2 - def setUp(self): - super().setUp() - gc.collect() - backend_empty_cache(torch_device) - - def tearDown(self): - super().tearDown() - gc.collect() - backend_empty_cache(torch_device) - def get_file_format(self, seed, shape): return f"gaussian_noise_s={seed}_shape={'_'.join([str(s) for s in shape])}.npy" @@ -84,18 +68,6 @@ class AutoencoderKLSingleFileTests(unittest.TestCase): assert numpy_cosine_similarity_distance(output_slice_1, output_slice_2) < 1e-4 - def test_single_file_components(self): - model = self.model_class.from_pretrained(self.repo_id) - model_single_file = self.model_class.from_single_file(self.ckpt_path, config=self.repo_id) - - PARAMS_TO_IGNORE = ["torch_dtype", "_name_or_path", "_use_default_values", "_diffusers_version"] - for param_name, param_value in model_single_file.config.items(): - if param_name in PARAMS_TO_IGNORE: - continue - assert model.config[param_name] == param_value, ( - f"{param_name} differs between pretrained loading and single file loading" - ) - def test_single_file_arguments(self): model_default = self.model_class.from_single_file(self.ckpt_path, config=self.repo_id) diff --git a/tests/single_file/test_model_wan_autoencoder_single_file.py b/tests/single_file/test_model_wan_autoencoder_single_file.py index a1f7155c10..0babf30234 100644 --- a/tests/single_file/test_model_wan_autoencoder_single_file.py +++ b/tests/single_file/test_model_wan_autoencoder_single_file.py @@ -13,50 +13,24 @@ # See the License for the specific language governing permissions and # limitations under the License. -import gc -import unittest from diffusers import ( AutoencoderKLWan, ) from ..testing_utils import ( - backend_empty_cache, enable_full_determinism, - require_torch_accelerator, - torch_device, ) +from .single_file_testing_utils import SingleFileModelTesterMixin enable_full_determinism() -@require_torch_accelerator -class AutoencoderKLWanSingleFileTests(unittest.TestCase): +class TestAutoencoderKLWanSingleFile(SingleFileModelTesterMixin): model_class = AutoencoderKLWan ckpt_path = ( "https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/vae/wan_2.1_vae.safetensors" ) repo_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" - - def setUp(self): - super().setUp() - gc.collect() - backend_empty_cache(torch_device) - - def tearDown(self): - super().tearDown() - gc.collect() - backend_empty_cache(torch_device) - - def test_single_file_components(self): - model = self.model_class.from_pretrained(self.repo_id, subfolder="vae") - model_single_file = self.model_class.from_single_file(self.ckpt_path) - - PARAMS_TO_IGNORE = ["torch_dtype", "_name_or_path", "_use_default_values", "_diffusers_version"] - for param_name, param_value in model_single_file.config.items(): - if param_name in PARAMS_TO_IGNORE: - continue - assert model.config[param_name] == param_value, ( - f"{param_name} differs between single file loading and pretrained loading" - ) + subfolder = "vae" diff --git a/tests/single_file/test_model_wan_transformer3d_single_file.py b/tests/single_file/test_model_wan_transformer3d_single_file.py index d7c758d3d9..b769092060 100644 --- a/tests/single_file/test_model_wan_transformer3d_single_file.py +++ b/tests/single_file/test_model_wan_transformer3d_single_file.py @@ -13,8 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import gc -import unittest import torch @@ -23,72 +21,26 @@ from diffusers import ( ) from ..testing_utils import ( - backend_empty_cache, enable_full_determinism, require_big_accelerator, - require_torch_accelerator, - torch_device, ) +from .single_file_testing_utils import SingleFileModelTesterMixin enable_full_determinism() -@require_torch_accelerator -class WanTransformer3DModelText2VideoSingleFileTest(unittest.TestCase): +class TestWanTransformer3DModelText2VideoSingleFile(SingleFileModelTesterMixin): model_class = WanTransformer3DModel ckpt_path = "https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/diffusion_models/wan2.1_t2v_1.3B_bf16.safetensors" repo_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" - - def setUp(self): - super().setUp() - gc.collect() - backend_empty_cache(torch_device) - - def tearDown(self): - super().tearDown() - gc.collect() - backend_empty_cache(torch_device) - - def test_single_file_components(self): - model = self.model_class.from_pretrained(self.repo_id, subfolder="transformer") - model_single_file = self.model_class.from_single_file(self.ckpt_path) - - PARAMS_TO_IGNORE = ["torch_dtype", "_name_or_path", "_use_default_values", "_diffusers_version"] - for param_name, param_value in model_single_file.config.items(): - if param_name in PARAMS_TO_IGNORE: - continue - assert model.config[param_name] == param_value, ( - f"{param_name} differs between single file loading and pretrained loading" - ) + subfolder = "transformer" @require_big_accelerator -@require_torch_accelerator -class WanTransformer3DModelImage2VideoSingleFileTest(unittest.TestCase): +class TestWanTransformer3DModelImage2VideoSingleFile(SingleFileModelTesterMixin): model_class = WanTransformer3DModel ckpt_path = "https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/diffusion_models/wan2.1_i2v_480p_14B_fp8_e4m3fn.safetensors" repo_id = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers" torch_dtype = torch.float8_e4m3fn - - def setUp(self): - super().setUp() - gc.collect() - backend_empty_cache(torch_device) - - def tearDown(self): - super().tearDown() - gc.collect() - backend_empty_cache(torch_device) - - def test_single_file_components(self): - model = self.model_class.from_pretrained(self.repo_id, subfolder="transformer", torch_dtype=self.torch_dtype) - model_single_file = self.model_class.from_single_file(self.ckpt_path, torch_dtype=self.torch_dtype) - - PARAMS_TO_IGNORE = ["torch_dtype", "_name_or_path", "_use_default_values", "_diffusers_version"] - for param_name, param_value in model_single_file.config.items(): - if param_name in PARAMS_TO_IGNORE: - continue - assert model.config[param_name] == param_value, ( - f"{param_name} differs between single file loading and pretrained loading" - ) + subfolder = "transformer" diff --git a/tests/single_file/test_sana_transformer.py b/tests/single_file/test_sana_transformer.py index c1543ba171..9e2adb93bf 100644 --- a/tests/single_file/test_sana_transformer.py +++ b/tests/single_file/test_sana_transformer.py @@ -1,23 +1,17 @@ -import gc -import unittest - from diffusers import ( SanaTransformer2DModel, ) from ..testing_utils import ( - backend_empty_cache, enable_full_determinism, - require_torch_accelerator, - torch_device, ) +from .single_file_testing_utils import SingleFileModelTesterMixin enable_full_determinism() -@require_torch_accelerator -class SanaTransformer2DModelSingleFileTests(unittest.TestCase): +class TestSanaTransformer2DModelSingleFile(SingleFileModelTesterMixin): model_class = SanaTransformer2DModel ckpt_path = ( "https://huggingface.co/Efficient-Large-Model/Sana_1600M_1024px/blob/main/checkpoints/Sana_1600M_1024px.pth" @@ -27,34 +21,4 @@ class SanaTransformer2DModelSingleFileTests(unittest.TestCase): ] repo_id = "Efficient-Large-Model/Sana_1600M_1024px_diffusers" - - def setUp(self): - super().setUp() - gc.collect() - backend_empty_cache(torch_device) - - def tearDown(self): - super().tearDown() - gc.collect() - backend_empty_cache(torch_device) - - def test_single_file_components(self): - model = self.model_class.from_pretrained(self.repo_id, subfolder="transformer") - model_single_file = self.model_class.from_single_file(self.ckpt_path) - - PARAMS_TO_IGNORE = ["torch_dtype", "_name_or_path", "_use_default_values", "_diffusers_version"] - for param_name, param_value in model_single_file.config.items(): - if param_name in PARAMS_TO_IGNORE: - continue - assert model.config[param_name] == param_value, ( - f"{param_name} differs between single file loading and pretrained loading" - ) - - def test_checkpoint_loading(self): - for ckpt_path in self.alternate_keys_ckpt_paths: - backend_empty_cache(torch_device) - model = self.model_class.from_single_file(ckpt_path) - - del model - gc.collect() - backend_empty_cache(torch_device) + subfolder = "transformer" diff --git a/tests/single_file/test_stable_diffusion_controlnet_img2img_single_file.py b/tests/single_file/test_stable_diffusion_controlnet_img2img_single_file.py index e558eeaf6f..141748b084 100644 --- a/tests/single_file/test_stable_diffusion_controlnet_img2img_single_file.py +++ b/tests/single_file/test_stable_diffusion_controlnet_img2img_single_file.py @@ -1,6 +1,5 @@ import gc import tempfile -import unittest import torch @@ -29,7 +28,7 @@ enable_full_determinism() @slow @require_torch_accelerator -class StableDiffusionControlNetPipelineSingleFileSlowTests(unittest.TestCase, SDSingleFileTesterMixin): +class TestStableDiffusionControlNetPipelineSingleFileSlow(SDSingleFileTesterMixin): pipeline_class = StableDiffusionControlNetPipeline ckpt_path = ( "https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5/blob/main/v1-5-pruned-emaonly.safetensors" @@ -39,13 +38,11 @@ class StableDiffusionControlNetPipelineSingleFileSlowTests(unittest.TestCase, SD ) repo_id = "stable-diffusion-v1-5/stable-diffusion-v1-5" - def setUp(self): - super().setUp() + def setup_method(self): gc.collect() backend_empty_cache(torch_device) - def tearDown(self): - super().tearDown() + def teardown_method(self): gc.collect() backend_empty_cache(torch_device) diff --git a/tests/single_file/test_stable_diffusion_controlnet_inpaint_single_file.py b/tests/single_file/test_stable_diffusion_controlnet_inpaint_single_file.py index 54224f51a9..8238866cbf 100644 --- a/tests/single_file/test_stable_diffusion_controlnet_inpaint_single_file.py +++ b/tests/single_file/test_stable_diffusion_controlnet_inpaint_single_file.py @@ -1,7 +1,7 @@ import gc import tempfile -import unittest +import pytest import torch from diffusers import ControlNetModel, StableDiffusionControlNetInpaintPipeline @@ -29,19 +29,17 @@ enable_full_determinism() @slow @require_torch_accelerator -class StableDiffusionControlNetInpaintPipelineSingleFileSlowTests(unittest.TestCase, SDSingleFileTesterMixin): +class TestStableDiffusionControlNetInpaintPipelineSingleFileSlow(SDSingleFileTesterMixin): pipeline_class = StableDiffusionControlNetInpaintPipeline ckpt_path = "https://huggingface.co/botp/stable-diffusion-v1-5-inpainting/blob/main/sd-v1-5-inpainting.ckpt" original_config = "https://raw.githubusercontent.com/runwayml/stable-diffusion/main/configs/stable-diffusion/v1-inpainting-inference.yaml" repo_id = "stable-diffusion-v1-5/stable-diffusion-inpainting" - def setUp(self): - super().setUp() + def setup_method(self): gc.collect() backend_empty_cache(torch_device) - def tearDown(self): - super().tearDown() + def teardown_method(self): gc.collect() backend_empty_cache(torch_device) @@ -115,7 +113,7 @@ class StableDiffusionControlNetInpaintPipelineSingleFileSlowTests(unittest.TestC super()._compare_component_configs(pipe, pipe_single_file) - @unittest.skip("runwayml original config repo does not exist") + @pytest.mark.skip(reason="runwayml original config repo does not exist") def test_single_file_components_with_original_config(self): controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_canny", variant="fp16") pipe = self.pipeline_class.from_pretrained(self.repo_id, controlnet=controlnet) @@ -125,7 +123,7 @@ class StableDiffusionControlNetInpaintPipelineSingleFileSlowTests(unittest.TestC super()._compare_component_configs(pipe, pipe_single_file) - @unittest.skip("runwayml original config repo does not exist") + @pytest.mark.skip(reason="runwayml original config repo does not exist") def test_single_file_components_with_original_config_local_files_only(self): controlnet = ControlNetModel.from_pretrained( "lllyasviel/control_v11p_sd15_canny", torch_dtype=torch.float16, variant="fp16" diff --git a/tests/single_file/test_stable_diffusion_controlnet_single_file.py b/tests/single_file/test_stable_diffusion_controlnet_single_file.py index e90e648a9d..80ef6c2574 100644 --- a/tests/single_file/test_stable_diffusion_controlnet_single_file.py +++ b/tests/single_file/test_stable_diffusion_controlnet_single_file.py @@ -1,6 +1,5 @@ import gc import tempfile -import unittest import torch @@ -29,7 +28,7 @@ enable_full_determinism() @slow @require_torch_accelerator -class StableDiffusionControlNetPipelineSingleFileSlowTests(unittest.TestCase, SDSingleFileTesterMixin): +class TestStableDiffusionControlNetPipelineSingleFileSlow(SDSingleFileTesterMixin): pipeline_class = StableDiffusionControlNetPipeline ckpt_path = ( "https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5/blob/main/v1-5-pruned-emaonly.safetensors" @@ -39,13 +38,11 @@ class StableDiffusionControlNetPipelineSingleFileSlowTests(unittest.TestCase, SD ) repo_id = "stable-diffusion-v1-5/stable-diffusion-v1-5" - def setUp(self): - super().setUp() + def setup_method(self): gc.collect() backend_empty_cache(torch_device) - def tearDown(self): - super().tearDown() + def teardown_method(self): gc.collect() backend_empty_cache(torch_device) diff --git a/tests/single_file/test_stable_diffusion_img2img_single_file.py b/tests/single_file/test_stable_diffusion_img2img_single_file.py index 387f09471d..e76846c800 100644 --- a/tests/single_file/test_stable_diffusion_img2img_single_file.py +++ b/tests/single_file/test_stable_diffusion_img2img_single_file.py @@ -1,5 +1,4 @@ import gc -import unittest import torch @@ -23,7 +22,7 @@ enable_full_determinism() @slow @require_torch_accelerator -class StableDiffusionImg2ImgPipelineSingleFileSlowTests(unittest.TestCase, SDSingleFileTesterMixin): +class TestStableDiffusionImg2ImgPipelineSingleFileSlow(SDSingleFileTesterMixin): pipeline_class = StableDiffusionImg2ImgPipeline ckpt_path = ( "https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5/blob/main/v1-5-pruned-emaonly.safetensors" @@ -33,13 +32,11 @@ class StableDiffusionImg2ImgPipelineSingleFileSlowTests(unittest.TestCase, SDSin ) repo_id = "stable-diffusion-v1-5/stable-diffusion-v1-5" - def setUp(self): - super().setUp() + def setup_method(self): gc.collect() backend_empty_cache(torch_device) - def tearDown(self): - super().tearDown() + def teardown_method(self): gc.collect() backend_empty_cache(torch_device) @@ -66,19 +63,17 @@ class StableDiffusionImg2ImgPipelineSingleFileSlowTests(unittest.TestCase, SDSin @slow @require_torch_accelerator -class StableDiffusion21Img2ImgPipelineSingleFileSlowTests(unittest.TestCase, SDSingleFileTesterMixin): +class TestStableDiffusion21Img2ImgPipelineSingleFileSlow(SDSingleFileTesterMixin): pipeline_class = StableDiffusionImg2ImgPipeline ckpt_path = "https://huggingface.co/stabilityai/stable-diffusion-2-1/blob/main/v2-1_768-ema-pruned.safetensors" original_config = "https://raw.githubusercontent.com/Stability-AI/stablediffusion/main/configs/stable-diffusion/v2-inference-v.yaml" repo_id = "stabilityai/stable-diffusion-2-1" - def setUp(self): - super().setUp() + def setup_method(self): gc.collect() backend_empty_cache(torch_device) - def tearDown(self): - super().tearDown() + def teardown_method(self): gc.collect() backend_empty_cache(torch_device) diff --git a/tests/single_file/test_stable_diffusion_inpaint_single_file.py b/tests/single_file/test_stable_diffusion_inpaint_single_file.py index 84636ec0f0..6e5d27cdff 100644 --- a/tests/single_file/test_stable_diffusion_inpaint_single_file.py +++ b/tests/single_file/test_stable_diffusion_inpaint_single_file.py @@ -1,6 +1,6 @@ import gc -import unittest +import pytest import torch from diffusers import ( @@ -23,19 +23,17 @@ enable_full_determinism() @slow @require_torch_accelerator -class StableDiffusionInpaintPipelineSingleFileSlowTests(unittest.TestCase, SDSingleFileTesterMixin): +class TestStableDiffusionInpaintPipelineSingleFileSlow(SDSingleFileTesterMixin): pipeline_class = StableDiffusionInpaintPipeline ckpt_path = "https://huggingface.co/botp/stable-diffusion-v1-5-inpainting/blob/main/sd-v1-5-inpainting.ckpt" original_config = "https://raw.githubusercontent.com/runwayml/stable-diffusion/main/configs/stable-diffusion/v1-inpainting-inference.yaml" repo_id = "botp/stable-diffusion-v1-5-inpainting" - def setUp(self): - super().setUp() + def setup_method(self): gc.collect() backend_empty_cache(torch_device) - def tearDown(self): - super().tearDown() + def teardown_method(self): gc.collect() backend_empty_cache(torch_device) @@ -70,18 +68,18 @@ class StableDiffusionInpaintPipelineSingleFileSlowTests(unittest.TestCase, SDSin assert pipe.unet.config.in_channels == 4 - @unittest.skip("runwayml original config has been removed") + @pytest.mark.skip(reason="runwayml original config has been removed") def test_single_file_components_with_original_config(self): return - @unittest.skip("runwayml original config has been removed") + @pytest.mark.skip(reason="runwayml original config has been removed") def test_single_file_components_with_original_config_local_files_only(self): return @slow @require_torch_accelerator -class StableDiffusion21InpaintPipelineSingleFileSlowTests(unittest.TestCase, SDSingleFileTesterMixin): +class TestStableDiffusion21InpaintPipelineSingleFileSlow(SDSingleFileTesterMixin): pipeline_class = StableDiffusionInpaintPipeline ckpt_path = ( "https://huggingface.co/stabilityai/stable-diffusion-2-inpainting/blob/main/512-inpainting-ema.safetensors" @@ -89,13 +87,11 @@ class StableDiffusion21InpaintPipelineSingleFileSlowTests(unittest.TestCase, SDS original_config = "https://raw.githubusercontent.com/Stability-AI/stablediffusion/main/configs/stable-diffusion/v2-inpainting-inference.yaml" repo_id = "stabilityai/stable-diffusion-2-inpainting" - def setUp(self): - super().setUp() + def setup_method(self): gc.collect() backend_empty_cache(torch_device) - def tearDown(self): - super().tearDown() + def teardown_method(self): gc.collect() backend_empty_cache(torch_device) diff --git a/tests/single_file/test_stable_diffusion_single_file.py b/tests/single_file/test_stable_diffusion_single_file.py index 4601b75c3a..377dedbc57 100644 --- a/tests/single_file/test_stable_diffusion_single_file.py +++ b/tests/single_file/test_stable_diffusion_single_file.py @@ -1,6 +1,5 @@ import gc import tempfile -import unittest import torch @@ -28,7 +27,7 @@ enable_full_determinism() @slow @require_torch_accelerator -class StableDiffusionPipelineSingleFileSlowTests(unittest.TestCase, SDSingleFileTesterMixin): +class TestStableDiffusionPipelineSingleFileSlow(SDSingleFileTesterMixin): pipeline_class = StableDiffusionPipeline ckpt_path = ( "https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5/blob/main/v1-5-pruned-emaonly.safetensors" @@ -38,13 +37,11 @@ class StableDiffusionPipelineSingleFileSlowTests(unittest.TestCase, SDSingleFile ) repo_id = "stable-diffusion-v1-5/stable-diffusion-v1-5" - def setUp(self): - super().setUp() + def setup_method(self): gc.collect() backend_empty_cache(torch_device) - def tearDown(self): - super().tearDown() + def teardown_method(self): gc.collect() backend_empty_cache(torch_device) @@ -90,19 +87,17 @@ class StableDiffusionPipelineSingleFileSlowTests(unittest.TestCase, SDSingleFile @slow -class StableDiffusion21PipelineSingleFileSlowTests(unittest.TestCase, SDSingleFileTesterMixin): +class TestStableDiffusion21PipelineSingleFileSlow(SDSingleFileTesterMixin): pipeline_class = StableDiffusionPipeline ckpt_path = "https://huggingface.co/stabilityai/stable-diffusion-2-1/blob/main/v2-1_768-ema-pruned.safetensors" original_config = "https://raw.githubusercontent.com/Stability-AI/stablediffusion/main/configs/stable-diffusion/v2-inference-v.yaml" repo_id = "stabilityai/stable-diffusion-2-1" - def setUp(self): - super().setUp() + def setup_method(self): gc.collect() backend_empty_cache(torch_device) - def tearDown(self): - super().tearDown() + def teardown_method(self): gc.collect() backend_empty_cache(torch_device) @@ -125,7 +120,7 @@ class StableDiffusion21PipelineSingleFileSlowTests(unittest.TestCase, SDSingleFi @nightly @slow @require_torch_accelerator -class StableDiffusionInstructPix2PixPipelineSingleFileSlowTests(unittest.TestCase, SDSingleFileTesterMixin): +class TestStableDiffusionInstructPix2PixPipelineSingleFileSlow(SDSingleFileTesterMixin): pipeline_class = StableDiffusionInstructPix2PixPipeline ckpt_path = "https://huggingface.co/timbrooks/instruct-pix2pix/blob/main/instruct-pix2pix-00-22000.safetensors" original_config = ( @@ -134,13 +129,11 @@ class StableDiffusionInstructPix2PixPipelineSingleFileSlowTests(unittest.TestCas repo_id = "timbrooks/instruct-pix2pix" single_file_kwargs = {"extract_ema": True} - def setUp(self): - super().setUp() + def setup_method(self): gc.collect() backend_empty_cache(torch_device) - def tearDown(self): - super().tearDown() + def teardown_method(self): gc.collect() backend_empty_cache(torch_device) diff --git a/tests/single_file/test_stable_diffusion_upscale_single_file.py b/tests/single_file/test_stable_diffusion_upscale_single_file.py index 39ec7b0194..ba4819fadf 100644 --- a/tests/single_file/test_stable_diffusion_upscale_single_file.py +++ b/tests/single_file/test_stable_diffusion_upscale_single_file.py @@ -1,5 +1,4 @@ import gc -import unittest import pytest import torch @@ -25,19 +24,17 @@ enable_full_determinism() @slow @require_torch_accelerator -class StableDiffusionUpscalePipelineSingleFileSlowTests(unittest.TestCase, SDSingleFileTesterMixin): +class TestStableDiffusionUpscalePipelineSingleFileSlow(SDSingleFileTesterMixin): pipeline_class = StableDiffusionUpscalePipeline ckpt_path = "https://huggingface.co/stabilityai/stable-diffusion-x4-upscaler/blob/main/x4-upscaler-ema.safetensors" original_config = "https://raw.githubusercontent.com/Stability-AI/stablediffusion/main/configs/stable-diffusion/x4-upscaling.yaml" repo_id = "stabilityai/stable-diffusion-x4-upscaler" - def setUp(self): - super().setUp() + def setup_method(self): gc.collect() backend_empty_cache(torch_device) - def tearDown(self): - super().tearDown() + def teardown_method(self): gc.collect() backend_empty_cache(torch_device) diff --git a/tests/single_file/test_stable_diffusion_xl_adapter_single_file.py b/tests/single_file/test_stable_diffusion_xl_adapter_single_file.py index 3de9ee7364..3d124fa8c2 100644 --- a/tests/single_file/test_stable_diffusion_xl_adapter_single_file.py +++ b/tests/single_file/test_stable_diffusion_xl_adapter_single_file.py @@ -1,6 +1,5 @@ import gc import tempfile -import unittest import torch @@ -32,7 +31,7 @@ enable_full_determinism() @slow @require_torch_accelerator -class StableDiffusionXLAdapterPipelineSingleFileSlowTests(unittest.TestCase, SDXLSingleFileTesterMixin): +class TestStableDiffusionXLAdapterPipelineSingleFileSlow(SDXLSingleFileTesterMixin): pipeline_class = StableDiffusionXLAdapterPipeline ckpt_path = "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_base_1.0.safetensors" repo_id = "stabilityai/stable-diffusion-xl-base-1.0" @@ -40,13 +39,11 @@ class StableDiffusionXLAdapterPipelineSingleFileSlowTests(unittest.TestCase, SDX "https://raw.githubusercontent.com/Stability-AI/generative-models/main/configs/inference/sd_xl_base.yaml" ) - def setUp(self): - super().setUp() + def setup_method(self): gc.collect() backend_empty_cache(torch_device) - def tearDown(self): - super().tearDown() + def teardown_method(self): gc.collect() backend_empty_cache(torch_device) diff --git a/tests/single_file/test_stable_diffusion_xl_controlnet_single_file.py b/tests/single_file/test_stable_diffusion_xl_controlnet_single_file.py index a0a1aba103..6f50370261 100644 --- a/tests/single_file/test_stable_diffusion_xl_controlnet_single_file.py +++ b/tests/single_file/test_stable_diffusion_xl_controlnet_single_file.py @@ -1,6 +1,5 @@ import gc import tempfile -import unittest import torch @@ -28,7 +27,7 @@ enable_full_determinism() @slow @require_torch_accelerator -class StableDiffusionXLControlNetPipelineSingleFileSlowTests(unittest.TestCase, SDXLSingleFileTesterMixin): +class TestStableDiffusionXLControlNetPipelineSingleFileSlow(SDXLSingleFileTesterMixin): pipeline_class = StableDiffusionXLControlNetPipeline ckpt_path = "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_base_1.0.safetensors" repo_id = "stabilityai/stable-diffusion-xl-base-1.0" @@ -36,13 +35,11 @@ class StableDiffusionXLControlNetPipelineSingleFileSlowTests(unittest.TestCase, "https://raw.githubusercontent.com/Stability-AI/generative-models/main/configs/inference/sd_xl_base.yaml" ) - def setUp(self): - super().setUp() + def setup_method(self): gc.collect() backend_empty_cache(torch_device) - def tearDown(self): - super().tearDown() + def teardown_method(self): gc.collect() backend_empty_cache(torch_device) diff --git a/tests/single_file/test_stable_diffusion_xl_img2img_single_file.py b/tests/single_file/test_stable_diffusion_xl_img2img_single_file.py index 810f412f8d..56657f37d9 100644 --- a/tests/single_file/test_stable_diffusion_xl_img2img_single_file.py +++ b/tests/single_file/test_stable_diffusion_xl_img2img_single_file.py @@ -1,5 +1,4 @@ import gc -import unittest import torch @@ -25,7 +24,7 @@ enable_full_determinism() @slow @require_torch_accelerator -class StableDiffusionXLImg2ImgPipelineSingleFileSlowTests(unittest.TestCase, SDXLSingleFileTesterMixin): +class TestStableDiffusionXLImg2ImgPipelineSingleFileSlow(SDXLSingleFileTesterMixin): pipeline_class = StableDiffusionXLImg2ImgPipeline ckpt_path = "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_base_1.0.safetensors" repo_id = "stabilityai/stable-diffusion-xl-base-1.0" @@ -33,13 +32,11 @@ class StableDiffusionXLImg2ImgPipelineSingleFileSlowTests(unittest.TestCase, SDX "https://raw.githubusercontent.com/Stability-AI/generative-models/main/configs/inference/sd_xl_base.yaml" ) - def setUp(self): - super().setUp() + def setup_method(self): gc.collect() backend_empty_cache(torch_device) - def tearDown(self): - super().tearDown() + def teardown_method(self): gc.collect() backend_empty_cache(torch_device) @@ -66,7 +63,7 @@ class StableDiffusionXLImg2ImgPipelineSingleFileSlowTests(unittest.TestCase, SDX @slow @require_torch_accelerator -class StableDiffusionXLImg2ImgRefinerPipelineSingleFileSlowTests(unittest.TestCase): +class StableDiffusionXLImg2ImgRefinerPipelineSingleFileSlowTests: pipeline_class = StableDiffusionXLImg2ImgPipeline ckpt_path = ( "https://huggingface.co/stabilityai/stable-diffusion-xl-refiner-1.0/blob/main/sd_xl_refiner_1.0.safetensors" diff --git a/tests/single_file/test_stable_diffusion_xl_instruct_pix2pix.py b/tests/single_file/test_stable_diffusion_xl_instruct_pix2pix.py index 011d59222a..d755b70105 100644 --- a/tests/single_file/test_stable_diffusion_xl_instruct_pix2pix.py +++ b/tests/single_file/test_stable_diffusion_xl_instruct_pix2pix.py @@ -1,5 +1,4 @@ import gc -import unittest import torch @@ -19,19 +18,17 @@ enable_full_determinism() @slow @require_torch_accelerator -class StableDiffusionXLInstructPix2PixPipeline(unittest.TestCase): +class StableDiffusionXLInstructPix2PixPipeline: pipeline_class = StableDiffusionXLInstructPix2PixPipeline ckpt_path = "https://huggingface.co/stabilityai/cosxl/blob/main/cosxl_edit.safetensors" original_config = None repo_id = "diffusers/sdxl-instructpix2pix-768" - def setUp(self): - super().setUp() + def setup_method(self): gc.collect() backend_empty_cache(torch_device) - def tearDown(self): - super().tearDown() + def teardown_method(self): gc.collect() backend_empty_cache(torch_device) diff --git a/tests/single_file/test_stable_diffusion_xl_single_file.py b/tests/single_file/test_stable_diffusion_xl_single_file.py index 0ad180de17..4e5319ca25 100644 --- a/tests/single_file/test_stable_diffusion_xl_single_file.py +++ b/tests/single_file/test_stable_diffusion_xl_single_file.py @@ -1,5 +1,4 @@ import gc -import unittest import torch @@ -22,7 +21,7 @@ enable_full_determinism() @slow @require_torch_accelerator -class StableDiffusionXLPipelineSingleFileSlowTests(unittest.TestCase, SDXLSingleFileTesterMixin): +class TestStableDiffusionXLPipelineSingleFileSlow(SDXLSingleFileTesterMixin): pipeline_class = StableDiffusionXLPipeline ckpt_path = "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_base_1.0.safetensors" repo_id = "stabilityai/stable-diffusion-xl-base-1.0" @@ -30,13 +29,11 @@ class StableDiffusionXLPipelineSingleFileSlowTests(unittest.TestCase, SDXLSingle "https://raw.githubusercontent.com/Stability-AI/generative-models/main/configs/inference/sd_xl_base.yaml" ) - def setUp(self): - super().setUp() + def setup_method(self): gc.collect() backend_empty_cache(torch_device) - def tearDown(self): - super().tearDown() + def teardown_method(self): gc.collect() backend_empty_cache(torch_device) From 0e12ba74542c6ecb02719ec3e5c6e993b85556e3 Mon Sep 17 00:00:00 2001 From: Yao Matrix Date: Tue, 30 Sep 2025 01:37:48 -0700 Subject: [PATCH 038/123] fix 3 xpu failures uts w/ latest pytorch (#12408) fix xpu ut failures w/ latest pytorch Signed-off-by: Yao, Matrix --- .../test_controlnet_hunyuandit.py | 2 +- tests/pipelines/flux/test_pipeline_flux.py | 11 ++-- tests/quantization/gguf/test_gguf.py | 50 +++++++++---------- 3 files changed, 34 insertions(+), 29 deletions(-) diff --git a/tests/pipelines/controlnet_hunyuandit/test_controlnet_hunyuandit.py b/tests/pipelines/controlnet_hunyuandit/test_controlnet_hunyuandit.py index 9619843779..bf31f2abcf 100644 --- a/tests/pipelines/controlnet_hunyuandit/test_controlnet_hunyuandit.py +++ b/tests/pipelines/controlnet_hunyuandit/test_controlnet_hunyuandit.py @@ -155,7 +155,7 @@ class HunyuanDiTControlNetPipelineFastTests(unittest.TestCase, PipelineTesterMix if torch_device == "xpu": expected_slice = np.array( - [0.6376953, 0.84375, 0.58691406, 0.48046875, 0.43652344, 0.5517578, 0.54248047, 0.5644531, 0.48217773] + [0.6948242, 0.89160156, 0.59375, 0.5078125, 0.57910156, 0.6035156, 0.58447266, 0.53564453, 0.52246094] ) else: expected_slice = np.array( diff --git a/tests/pipelines/flux/test_pipeline_flux.py b/tests/pipelines/flux/test_pipeline_flux.py index c3e8517d64..1ddbd4ba3d 100644 --- a/tests/pipelines/flux/test_pipeline_flux.py +++ b/tests/pipelines/flux/test_pipeline_flux.py @@ -15,6 +15,7 @@ from diffusers import ( ) from ...testing_utils import ( + Expectations, backend_empty_cache, nightly, numpy_cosine_similarity_distance, @@ -276,10 +277,14 @@ class FluxPipelineSlowTests(unittest.TestCase): image = pipe(**inputs).images[0] image_slice = image[0, :10, :10] # fmt: off - expected_slice = np.array( - [0.3242, 0.3203, 0.3164, 0.3164, 0.3125, 0.3125, 0.3281, 0.3242, 0.3203, 0.3301, 0.3262, 0.3242, 0.3281, 0.3242, 0.3203, 0.3262, 0.3262, 0.3164, 0.3262, 0.3281, 0.3184, 0.3281, 0.3281, 0.3203, 0.3281, 0.3281, 0.3164, 0.3320, 0.3320, 0.3203], - dtype=np.float32, + + expected_slices = Expectations( + { + ("cuda", None): np.array([0.3242, 0.3203, 0.3164, 0.3164, 0.3125, 0.3125, 0.3281, 0.3242, 0.3203, 0.3301, 0.3262, 0.3242, 0.3281, 0.3242, 0.3203, 0.3262, 0.3262, 0.3164, 0.3262, 0.3281, 0.3184, 0.3281, 0.3281, 0.3203, 0.3281, 0.3281, 0.3164, 0.3320, 0.3320, 0.3203], dtype=np.float32,), + ("xpu", 3): np.array([0.3301, 0.3281, 0.3359, 0.3203, 0.3203, 0.3281, 0.3281, 0.3301, 0.3340, 0.3281, 0.3320, 0.3359, 0.3281, 0.3301, 0.3320, 0.3242, 0.3301, 0.3281, 0.3242, 0.3320, 0.3320, 0.3281, 0.3320, 0.3320, 0.3262, 0.3320, 0.3301, 0.3301, 0.3359, 0.3320], dtype=np.float32,), + } ) + expected_slice = expected_slices.get_expectation() # fmt: on max_diff = numpy_cosine_similarity_distance(expected_slice.flatten(), image_slice.flatten()) diff --git a/tests/quantization/gguf/test_gguf.py b/tests/quantization/gguf/test_gguf.py index 38322459e7..0f4fd408a7 100644 --- a/tests/quantization/gguf/test_gguf.py +++ b/tests/quantization/gguf/test_gguf.py @@ -360,33 +360,33 @@ class SD35LargeGGUFSingleFileTests(GGUFSingleFileTesterMixin, unittest.TestCase) { ("xpu", 3): np.array( [ - 0.1953125, - 0.3125, - 0.31445312, - 0.13085938, - 0.30664062, - 0.29296875, - 0.11523438, - 0.2890625, + 0.16796875, + 0.27929688, 0.28320312, - 0.16601562, - 0.3046875, - 0.328125, - 0.140625, - 0.31640625, + 0.11328125, + 0.27539062, + 0.26171875, + 0.10742188, + 0.26367188, + 0.26171875, + 0.1484375, + 0.2734375, + 0.296875, + 0.13476562, + 0.2890625, + 0.30078125, + 0.1171875, + 0.28125, + 0.28125, + 0.16015625, + 0.31445312, + 0.30078125, + 0.15625, 0.32421875, - 0.12304688, - 0.3046875, - 0.3046875, - 0.17578125, - 0.3359375, - 0.3203125, - 0.16601562, - 0.34375, - 0.31640625, - 0.15429688, - 0.328125, - 0.31054688, + 0.296875, + 0.14453125, + 0.30859375, + 0.2890625, ] ), ("cuda", 7): np.array( From b59654544bbaf5c6040e670397351abe0e543a75 Mon Sep 17 00:00:00 2001 From: Lucain Date: Tue, 30 Sep 2025 13:32:33 +0200 Subject: [PATCH 039/123] Install latest prerelease from huggingface_hub when installing transformers from main (#12395) * Allow prerelease when installing transformers from main * maybe better * maybe better * and now? * just bored * should be better * works now --- .github/workflows/pr_modular_tests.yml | 5 ++--- .github/workflows/pr_tests.yml | 12 +++++------- .github/workflows/pr_tests_gpu.yml | 13 +++++-------- tests/models/test_modeling_common.py | 8 +++++--- 4 files changed, 17 insertions(+), 21 deletions(-) diff --git a/.github/workflows/pr_modular_tests.yml b/.github/workflows/pr_modular_tests.yml index c6e87e642d..75258771e4 100644 --- a/.github/workflows/pr_modular_tests.yml +++ b/.github/workflows/pr_modular_tests.yml @@ -110,9 +110,8 @@ jobs: run: | python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" python -m uv pip install -e [quality,test] - # Stopping this update temporarily until the Hub RC is fully shipped and integrated. - # pip uninstall transformers -y && python -m uv pip install -U transformers@git+https://github.com/huggingface/transformers.git --no-deps - # pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git --no-deps + pip uninstall transformers -y && pip uninstall huggingface_hub -y && python -m uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git + pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git --no-deps - name: Environment run: | diff --git a/.github/workflows/pr_tests.yml b/.github/workflows/pr_tests.yml index ebfe9f442f..1543b264b0 100644 --- a/.github/workflows/pr_tests.yml +++ b/.github/workflows/pr_tests.yml @@ -116,9 +116,8 @@ jobs: run: | python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" python -m uv pip install -e [quality,test] - # Stopping this update temporarily until the Hub RC is fully shipped and integrated. - # pip uninstall transformers -y && python -m uv pip install -U transformers@git+https://github.com/huggingface/transformers.git --no-deps - # pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git --no-deps + pip uninstall transformers -y && pip uninstall huggingface_hub -y && python -m uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git + pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git --no-deps - name: Environment run: | @@ -254,10 +253,9 @@ jobs: python -m uv pip install -e [quality,test] # TODO (sayakpaul, DN6): revisit `--no-deps` python -m pip install -U peft@git+https://github.com/huggingface/peft.git --no-deps - # Stopping this update temporarily until the Hub RC is fully shipped and integrated. - # python -m uv pip install -U transformers@git+https://github.com/huggingface/transformers.git --no-deps - # python -m uv pip install -U tokenizers - # pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git --no-deps + python -m uv pip install -U tokenizers + pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git --no-deps + pip uninstall transformers -y && pip uninstall huggingface_hub -y && python -m uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git - name: Environment run: | diff --git a/.github/workflows/pr_tests_gpu.yml b/.github/workflows/pr_tests_gpu.yml index 1a8d5f6b81..89b6abe20d 100644 --- a/.github/workflows/pr_tests_gpu.yml +++ b/.github/workflows/pr_tests_gpu.yml @@ -132,9 +132,8 @@ jobs: run: | python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" python -m uv pip install -e [quality,test] - # Stopping this update temporarily until the Hub RC is fully shipped and integrated. - # pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git - # pip uninstall transformers -y && python -m uv pip install -U transformers@git+https://github.com/huggingface/transformers.git --no-deps + pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git + pip uninstall transformers -y && pip uninstall huggingface_hub -y && python -m uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git - name: Environment run: | @@ -204,9 +203,8 @@ jobs: python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" python -m uv pip install -e [quality,test] python -m uv pip install peft@git+https://github.com/huggingface/peft.git - # Stopping this update temporarily until the Hub RC is fully shipped and integrated. - # pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git - # pip uninstall transformers -y && python -m uv pip install -U transformers@git+https://github.com/huggingface/transformers.git --no-deps + pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git + pip uninstall transformers -y && pip uninstall huggingface_hub -y && python -m uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git - name: Environment run: | @@ -268,8 +266,7 @@ jobs: - name: Install dependencies run: | python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - # Stopping this update temporarily until the Hub RC is fully shipped and integrated. - # pip uninstall transformers -y && python -m uv pip install -U transformers@git+https://github.com/huggingface/transformers.git --no-deps + pip uninstall transformers -y && pip uninstall huggingface_hub -y && python -m uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git python -m uv pip install -e [quality,test,training] - name: Environment diff --git a/tests/models/test_modeling_common.py b/tests/models/test_modeling_common.py index 90ded6a7ec..9b1c6b50dc 100644 --- a/tests/models/test_modeling_common.py +++ b/tests/models/test_modeling_common.py @@ -243,8 +243,8 @@ class ModelUtilsTest(unittest.TestCase): else: _ = load_model(repo_id) - warning_message = str(warning.warnings[0].message) - self.assertIn("This serialization format is now deprecated to standardize the serialization", warning_message) + warning_messages = " ".join(str(w.message) for w in warning.warnings) + self.assertIn("This serialization format is now deprecated to standardize the serialization", warning_messages) # Local tests are already covered down below. @parameterized.expand( @@ -298,11 +298,13 @@ class ModelUtilsTest(unittest.TestCase): raise_for_status=mock.Mock(side_effect=HfHubHTTPError("Server down", response=mock.Mock())), json=mock.Mock(return_value={}), ) + client_mock = mock.Mock() + client_mock.get.return_value = error_response with tempfile.TemporaryDirectory() as tmpdir: model = FluxTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=tmpdir) - with mock.patch("requests.Session.get", return_value=error_response): + with mock.patch("huggingface_hub.hf_api.get_session", return_value=client_mock): # Should fail with local_files_only=False (network required) # We would make a network call with model_info with self.assertRaises(OSError): From d7a1a0363feb7694568695394a7edfdd44a18bca Mon Sep 17 00:00:00 2001 From: Steven Liu <59462357+stevhliu@users.noreply.github.com> Date: Tue, 30 Sep 2025 09:33:41 -0700 Subject: [PATCH 040/123] [docs] CP (#12331) * init * feedback * feedback * feedback * feedback * feedback * feedback --- docs/source/en/_toctree.yml | 4 +- docs/source/en/api/parallel.md | 2 +- .../en/training/distributed_inference.md | 64 +++++++++++++++++-- 3 files changed, 63 insertions(+), 7 deletions(-) diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index ada5e38895..fb4fdf2098 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -70,8 +70,6 @@ title: Reduce memory usage - local: optimization/speed-memory-optims title: Compiling and offloading quantized models - - local: api/parallel - title: Parallel inference - title: Community optimizations sections: - local: optimization/pruna @@ -282,6 +280,8 @@ title: Outputs - local: api/quantization title: Quantization + - local: api/parallel + title: Parallel inference - title: Modular sections: - local: api/modular_diffusers/pipeline diff --git a/docs/source/en/api/parallel.md b/docs/source/en/api/parallel.md index e38ffe571e..f2a6bee391 100644 --- a/docs/source/en/api/parallel.md +++ b/docs/source/en/api/parallel.md @@ -11,7 +11,7 @@ specific language governing permissions and limitations under the License. --> # Parallelism -Parallelism strategies help speed up diffusion transformers by distributing computations across multiple devices, allowing for faster inference/training times. +Parallelism strategies help speed up diffusion transformers by distributing computations across multiple devices, allowing for faster inference/training times. Refer to the [Distributed inferece](../training/distributed_inference) guide to learn more. ## ParallelConfig diff --git a/docs/source/en/training/distributed_inference.md b/docs/source/en/training/distributed_inference.md index 58ec77f75b..586f765709 100644 --- a/docs/source/en/training/distributed_inference.md +++ b/docs/source/en/training/distributed_inference.md @@ -226,8 +226,64 @@ with torch.no_grad(): image[0].save("split_transformer.png") ``` -## Resources +By selectively loading and unloading the models you need at a given stage and sharding the largest models across multiple GPUs, it is possible to run inference with large models on consumer GPUs. -- Take a look at this [script](https://gist.github.com/sayakpaul/cfaebd221820d7b43fae638b4dfa01ba) for a minimal example of distributed inference with Accelerate. -- For more details, check out Accelerate's [Distributed inference](https://huggingface.co/docs/accelerate/en/usage_guides/distributed_inference#distributed-inference-with-accelerate) guide. -- The `device_map` argument assign models or an entire pipeline to devices. Refer to the [device placement](../using-diffusers/loading#device-placement) docs for more information. \ No newline at end of file +## Context parallelism + +[Context parallelism](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=context_parallelism) splits input sequences across multiple GPUs to reduce memory usage. Each GPU processes its own slice of the sequence. + +Use [`~ModelMixin.set_attention_backend`] to switch to a more optimized attention backend. Refer to this [table](../optimization/attention_backends#available-backends) for a complete list of available backends. + +### Ring Attention + +Key (K) and value (V) representations communicate between devices using [Ring Attention](https://huggingface.co/papers/2310.01889). This ensures each split sees every other token's K/V. Each GPU computes attention for its local K/V and passes it to the next GPU in the ring. No single GPU holds the full sequence, which reduces communication latency. + +Pass a [`ContextParallelConfig`] to the `parallel_config` argument of the transformer model. The config supports the `ring_degree` argument that determines how many devices to use for Ring Attention. + +```py +import torch +from diffusers import AutoModel, QwenImagePipeline, ContextParallelConfig + +try: + torch.distributed.init_process_group("nccl") + rank = torch.distributed.get_rank() + device = torch.device("cuda", rank % torch.cuda.device_count()) + torch.cuda.set_device(device) + + transformer = AutoModel.from_pretrained("Qwen/Qwen-Image", subfolder="transformer", torch_dtype=torch.bfloat16, parallel_config=ContextParallelConfig(ring_degree=2)) + pipeline = QwenImagePipeline.from_pretrained("Qwen/Qwen-Image", transformer=transformer, torch_dtype=torch.bfloat16, device_map="cuda") + pipeline.transformer.set_attention_backend("flash") + + prompt = """ + cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California + highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain + """ + + # Must specify generator so all ranks start with same latents (or pass your own) + generator = torch.Generator().manual_seed(42) + image = pipeline(prompt, num_inference_steps=50, generator=generator).images[0] + + if rank == 0: + image.save("output.png") + +except Exception as e: + print(f"An error occurred: {e}") + torch.distributed.breakpoint() + raise + +finally: + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() +``` + +### Ulysses Attention + +[Ulysses Attention](https://huggingface.co/papers/2309.14509) splits a sequence across GPUs and performs an *all-to-all* communication (every device sends/receives data to every other device). Each GPU ends up with all tokens for only a subset of attention heads. Each GPU computes attention locally on all tokens for its head, then performs another all-to-all to regroup results by tokens for the next layer. + +[`ContextParallelConfig`] supports Ulysses Attention through the `ulysses_degree` argument. This determines how many devices to use for Ulysses Attention. + +Pass the [`ContextParallelConfig`] to [`~ModelMixin.enable_parallelism`]. + +```py +pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2)) +``` \ No newline at end of file From cc5b31ffc985efff096735ec290876617e0cabbb Mon Sep 17 00:00:00 2001 From: Steven Liu <59462357+stevhliu@users.noreply.github.com> Date: Tue, 30 Sep 2025 10:11:19 -0700 Subject: [PATCH 041/123] [docs] Migrate syntax (#12390) * change syntax * make style --- docs/source/en/api/configuration.md | 7 +- docs/source/en/api/loaders/ip_adapter.md | 7 +- docs/source/en/api/loaders/lora.md | 7 +- docs/source/en/api/loaders/peft.md | 7 +- .../en/api/loaders/textual_inversion.md | 7 +- docs/source/en/api/loaders/transformer_sd3.md | 7 +- docs/source/en/api/loaders/unet.md | 7 +- .../en/api/models/consistency_decoder_vae.md | 7 +- docs/source/en/api/models/transformer2d.md | 7 +- docs/source/en/api/outputs.md | 7 +- docs/source/en/api/pipelines/allegro.md | 7 +- docs/source/en/api/pipelines/animatediff.md | 21 +- .../en/api/pipelines/attend_and_excite.md | 7 +- docs/source/en/api/pipelines/audioldm.md | 7 +- docs/source/en/api/pipelines/audioldm2.md | 7 +- docs/source/en/api/pipelines/aura_flow.md | 7 +- .../source/en/api/pipelines/blip_diffusion.md | 7 +- docs/source/en/api/pipelines/chroma.md | 7 +- docs/source/en/api/pipelines/cogview3.md | 7 +- docs/source/en/api/pipelines/cogview4.md | 7 +- docs/source/en/api/pipelines/consisid.md | 7 +- .../en/api/pipelines/control_flux_inpaint.md | 7 +- docs/source/en/api/pipelines/controlnet.md | 7 +- .../en/api/pipelines/controlnet_flux.md | 7 +- .../en/api/pipelines/controlnet_hunyuandit.md | 7 +- .../source/en/api/pipelines/controlnet_sd3.md | 7 +- .../en/api/pipelines/controlnet_sdxl.md | 14 +- docs/source/en/api/pipelines/controlnetxs.md | 7 +- .../en/api/pipelines/controlnetxs_sdxl.md | 14 +- docs/source/en/api/pipelines/cosmos.md | 7 +- .../en/api/pipelines/dance_diffusion.md | 7 +- docs/source/en/api/pipelines/ddpm.md | 7 +- docs/source/en/api/pipelines/dit.md | 7 +- docs/source/en/api/pipelines/flux.md | 23 +- docs/source/en/api/pipelines/framepack.md | 7 +- docs/source/en/api/pipelines/hidream.md | 7 +- docs/source/en/api/pipelines/hunyuandit.md | 14 +- docs/source/en/api/pipelines/i2vgenxl.md | 7 +- docs/source/en/api/pipelines/kandinsky.md | 14 +- docs/source/en/api/pipelines/kandinsky3.md | 14 +- docs/source/en/api/pipelines/kandinsky_v22.md | 14 +- docs/source/en/api/pipelines/kolors.md | 14 +- .../en/api/pipelines/latent_diffusion.md | 7 +- docs/source/en/api/pipelines/latte.md | 7 +- docs/source/en/api/pipelines/ledits_pp.md | 14 +- docs/source/en/api/pipelines/lumina.md | 7 +- docs/source/en/api/pipelines/lumina2.md | 7 +- docs/source/en/api/pipelines/marigold.md | 47 +- docs/source/en/api/pipelines/mochi.md | 19 +- docs/source/en/api/pipelines/musicldm.md | 7 +- docs/source/en/api/pipelines/omnigen.md | 7 +- docs/source/en/api/pipelines/overview.md | 15 +- docs/source/en/api/pipelines/pag.md | 7 +- .../en/api/pipelines/paint_by_example.md | 7 +- docs/source/en/api/pipelines/panorama.md | 7 +- docs/source/en/api/pipelines/pia.md | 14 +- docs/source/en/api/pipelines/pix2pix.md | 7 +- docs/source/en/api/pipelines/pixart.md | 21 +- docs/source/en/api/pipelines/pixart_sigma.md | 28 +- docs/source/en/api/pipelines/qwenimage.md | 14 +- docs/source/en/api/pipelines/sana.md | 14 +- docs/source/en/api/pipelines/sana_sprint.md | 7 +- .../api/pipelines/self_attention_guidance.md | 7 +- .../pipelines/semantic_stable_diffusion.md | 7 +- docs/source/en/api/pipelines/shap_e.md | 7 +- .../source/en/api/pipelines/stable_cascade.md | 15 +- .../pipelines/stable_diffusion/depth2img.md | 11 +- .../api/pipelines/stable_diffusion/gligen.md | 11 +- .../stable_diffusion/image_variation.md | 7 +- .../api/pipelines/stable_diffusion/img2img.md | 7 +- .../api/pipelines/stable_diffusion/inpaint.md | 11 +- .../stable_diffusion/latent_upscale.md | 11 +- .../stable_diffusion/ldm3d_diffusion.md | 7 +- .../pipelines/stable_diffusion/sdxl_turbo.md | 11 +- .../stable_diffusion/stable_diffusion_2.md | 11 +- .../stable_diffusion/stable_diffusion_3.md | 21 +- .../stable_diffusion/stable_diffusion_safe.md | 7 +- .../stable_diffusion/stable_diffusion_xl.md | 11 +- .../en/api/pipelines/stable_diffusion/svd.md | 15 +- .../pipelines/stable_diffusion/text2img.md | 11 +- .../api/pipelines/stable_diffusion/upscale.md | 11 +- docs/source/en/api/pipelines/stable_unclip.md | 14 +- docs/source/en/api/pipelines/text_to_video.md | 7 +- .../en/api/pipelines/text_to_video_zero.md | 7 +- docs/source/en/api/pipelines/unclip.md | 7 +- docs/source/en/api/pipelines/unidiffuser.md | 14 +- .../en/api/pipelines/value_guided_sampling.md | 14 +- docs/source/en/api/quantization.md | 7 +- docs/source/en/api/schedulers/ddim.md | 7 +- docs/source/en/api/schedulers/score_sde_vp.md | 7 +- docs/source/en/conceptual/evaluation.md | 58 +- docs/source/en/optimization/coreml.md | 7 +- docs/source/en/optimization/fp16.md | 7 +- docs/source/en/optimization/mps.md | 7 +- docs/source/en/optimization/neuron.md | 7 +- docs/source/en/optimization/onnx.md | 7 +- docs/source/en/optimization/xformers.md | 14 +- docs/source/en/quantization/bitsandbytes.md | 21 +- docs/source/en/training/controlnet.md | 21 +- docs/source/en/training/create_dataset.md | 14 +- docs/source/en/training/custom_diffusion.md | 38 +- .../en/training/distributed_inference.md | 3 + docs/source/en/training/dreambooth.md | 64 +- docs/source/en/training/instructpix2pix.md | 29 +- docs/source/en/training/kandinsky.md | 35 +- docs/source/en/training/lcm_distill.md | 14 +- docs/source/en/training/lora.md | 35 +- docs/source/en/training/sdxl.md | 28 +- docs/source/en/training/t2i_adapters.md | 21 +- docs/source/en/training/text2image.md | 28 +- docs/source/en/training/text_inversion.md | 21 +- .../en/training/unconditional_training.md | 21 +- docs/source/en/training/wuerstchen.md | 21 +- docs/source/en/tutorials/basic_training.md | 21 +- .../conditional_image_generation.md | 21 +- .../using-diffusers/controlling_generation.md | 20 +- docs/source/en/using-diffusers/diffedit.md | 7 +- docs/source/en/using-diffusers/img2img.md | 21 +- .../inference_with_tcd_lora.md | 5 +- docs/source/en/using-diffusers/inpaint.md | 14 +- docs/source/en/using-diffusers/kandinsky.md | 47 +- docs/source/en/using-diffusers/pag.md | 7 +- docs/source/en/using-diffusers/sdxl.md | 36 +- docs/source/en/using-diffusers/shap-e.md | 7 +- .../unconditional_image_generation.md | 7 +- .../en/using-diffusers/weighted_prompts.md | 7 +- .../en/using-diffusers/write_own_pipeline.md | 29 +- docs/source/ja/installation.md | 7 +- docs/source/ja/quicktour.md | 28 +- docs/source/ja/stable_diffusion.md | 14 +- docs/source/ja/tutorials/autopipeline.md | 7 +- .../stable_diffusion/stable_diffusion_xl.md | 14 +- docs/source/ko/conceptual/evaluation.md | 44 +- docs/source/ko/installation.md | 7 +- docs/source/ko/optimization/coreml.md | 7 +- docs/source/ko/optimization/fp16.md | 22 +- docs/source/ko/optimization/mps.md | 7 +- docs/source/ko/optimization/xformers.md | 14 +- docs/source/ko/quicktour.md | 28 +- docs/source/ko/stable_diffusion.md | 14 +- docs/source/ko/training/controlnet.md | 7 +- docs/source/ko/training/create_dataset.md | 14 +- .../ko/training/distributed_inference.md | 5 +- docs/source/ko/training/dreambooth.md | 14 +- docs/source/ko/training/lora.md | 21 +- docs/source/ko/training/text2image.md | 7 +- docs/source/ko/training/text_inversion.md | 36 +- .../ko/training/unconditional_training.md | 7 +- docs/source/ko/tutorials/basic_training.md | 7 +- .../using-diffusers/controlling_generation.md | 20 +- .../custom_pipeline_overview.md | 7 +- docs/source/ko/using-diffusers/diffedit.md | 7 +- docs/source/ko/using-diffusers/img2img.md | 7 +- docs/source/ko/using-diffusers/inpaint.md | 7 +- docs/source/ko/using-diffusers/kandinsky.md | 47 +- docs/source/ko/using-diffusers/loading.md | 14 +- .../ko/using-diffusers/loading_adapters.md | 39 +- .../ko/using-diffusers/other-formats.md | 7 +- docs/source/ko/using-diffusers/schedulers.md | 13 +- docs/source/ko/using-diffusers/shap-e.md | 7 +- .../unconditional_image_generation.md | 7 +- .../ko/using-diffusers/write_own_pipeline.md | 29 +- docs/source/pt/installation.md | 7 +- docs/source/pt/quicktour.md | 28 +- docs/source/zh/conceptual/evaluation.md | 44 +- docs/source/zh/installation.md | 7 +- docs/source/zh/optimization/coreml.md | 7 +- docs/source/zh/optimization/fp16.md | 7 +- docs/source/zh/optimization/mps.md | 7 +- docs/source/zh/optimization/neuron.md | 7 +- docs/source/zh/optimization/onnx.md | 7 +- docs/source/zh/optimization/xformers.md | 14 +- docs/source/zh/quicktour.md | 31 +- docs/source/zh/stable_diffusion.md | 522 +++++++------- docs/source/zh/training/controlnet.md | 28 +- .../zh/training/distributed_inference.md | 7 +- docs/source/zh/training/dreambooth.md | 64 +- docs/source/zh/training/instructpix2pix.md | 29 +- docs/source/zh/training/kandinsky.md | 39 +- docs/source/zh/training/lora.md | 35 +- docs/source/zh/training/text2image.md | 35 +- docs/source/zh/training/text_inversion.md | 21 +- docs/source/zh/training/wuerstchen.md | 21 +- examples/community/matryoshka.py | 21 +- .../pipeline_stable_diffusion_boxdiff.py | 14 +- .../pipeline_stable_diffusion_pag.py | 10 +- examples/model_search/pipeline_easy.py | 54 +- src/diffusers/guiders/guider_utils.py | 10 +- src/diffusers/loaders/lora_base.py | 12 +- src/diffusers/loaders/lora_pipeline.py | 36 +- src/diffusers/models/attention.py | 6 +- src/diffusers/models/attention_processor.py | 6 +- src/diffusers/models/auto_model.py | 10 +- .../models/autoencoders/autoencoder_kl.py | 12 +- .../models/controlnets/controlnet_sd3.py | 12 +- .../models/controlnets/controlnet_xs.py | 12 +- src/diffusers/models/modeling_flax_utils.py | 12 +- src/diffusers/models/modeling_utils.py | 18 +- .../transformers/auraflow_transformer_2d.py | 12 +- .../transformers/cogvideox_transformer_3d.py | 12 +- .../transformers/hunyuan_transformer_2d.py | 12 +- .../transformers/pixart_transformer_2d.py | 12 +- .../models/transformers/transformer_sd3.py | 12 +- .../models/unets/unet_2d_condition.py | 12 +- .../models/unets/unet_3d_condition.py | 12 +- src/diffusers/models/unets/unet_i2vgen_xl.py | 12 +- .../models/unets/unet_motion_model.py | 12 +- .../modular_pipelines/components_manager.py | 6 +- .../flux/modular_pipeline.py | 6 +- .../modular_pipelines/modular_pipeline.py | 38 +- src/diffusers/modular_pipelines/node_utils.py | 661 ++++++++++++++++++ .../qwenimage/modular_pipeline.py | 12 +- .../stable_diffusion_xl/modular_pipeline.py | 6 +- .../modular_pipelines/wan/modular_pipeline.py | 6 +- src/diffusers/pipelines/auto_pipeline.py | 24 +- .../controlnet/pipeline_controlnet_inpaint.py | 13 +- .../controlnet/pipeline_flax_controlnet.py | 8 +- .../versatile_diffusion/modeling_text_unet.py | 12 +- .../pag/pipeline_pag_controlnet_sd_inpaint.py | 14 +- .../pipeline_paint_by_example.py | 6 +- .../pipelines/pipeline_flax_utils.py | 8 +- src/diffusers/pipelines/pipeline_utils.py | 67 +- .../pipeline_flax_stable_diffusion.py | 8 +- .../pipeline_flax_stable_diffusion_img2img.py | 8 +- .../pipeline_flax_stable_diffusion_inpaint.py | 14 +- .../pipeline_stable_diffusion_diffedit.py | 6 +- .../pipeline_stable_diffusion_k_diffusion.py | 6 +- .../deprecated/scheduling_karras_ve.py | 10 +- .../scheduling_consistency_models.py | 6 +- .../scheduling_cosine_dpmsolver_multistep.py | 8 +- .../scheduling_dpmsolver_multistep.py | 8 +- .../scheduling_dpmsolver_multistep_inverse.py | 8 +- .../scheduling_dpmsolver_singlestep.py | 8 +- .../scheduling_edm_dpmsolver_multistep.py | 8 +- .../schedulers/scheduling_sasolver.py | 8 +- src/diffusers/schedulers/scheduling_utils.py | 10 +- .../schedulers/scheduling_utils_flax.py | 17 +- src/diffusers/utils/dynamic_modules_utils.py | 24 +- src/diffusers/utils/outputs.py | 8 +- 239 files changed, 1947 insertions(+), 2656 deletions(-) create mode 100644 src/diffusers/modular_pipelines/node_utils.py diff --git a/docs/source/en/api/configuration.md b/docs/source/en/api/configuration.md index bc58e190b8..328e109e1e 100644 --- a/docs/source/en/api/configuration.md +++ b/docs/source/en/api/configuration.md @@ -14,11 +14,8 @@ specific language governing permissions and limitations under the License. Schedulers from [`~schedulers.scheduling_utils.SchedulerMixin`] and models from [`ModelMixin`] inherit from [`ConfigMixin`] which stores all the parameters that are passed to their respective `__init__` methods in a JSON-configuration file. - - -To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with `hf auth login`. - - +> [!TIP] +> To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with `hf auth login`. ## ConfigMixin diff --git a/docs/source/en/api/loaders/ip_adapter.md b/docs/source/en/api/loaders/ip_adapter.md index 0c94bcb220..c2c45bee10 100644 --- a/docs/source/en/api/loaders/ip_adapter.md +++ b/docs/source/en/api/loaders/ip_adapter.md @@ -14,11 +14,8 @@ specific language governing permissions and limitations under the License. [IP-Adapter](https://hf.co/papers/2308.06721) is a lightweight adapter that enables prompting a diffusion model with an image. This method decouples the cross-attention layers of the image and text features. The image features are generated from an image encoder. - - -Learn how to load an IP-Adapter checkpoint and image in the IP-Adapter [loading](../../using-diffusers/loading_adapters#ip-adapter) guide, and you can see how to use it in the [usage](../../using-diffusers/ip_adapter) guide. - - +> [!TIP] +> Learn how to load an IP-Adapter checkpoint and image in the IP-Adapter [loading](../../using-diffusers/loading_adapters#ip-adapter) guide, and you can see how to use it in the [usage](../../using-diffusers/ip_adapter) guide. ## IPAdapterMixin diff --git a/docs/source/en/api/loaders/lora.md b/docs/source/en/api/loaders/lora.md index da5c3842c6..bf22a32d74 100644 --- a/docs/source/en/api/loaders/lora.md +++ b/docs/source/en/api/loaders/lora.md @@ -33,11 +33,8 @@ LoRA is a fast and lightweight training method that inserts and trains a signifi - [`QwenImageLoraLoaderMixin`] provides similar functions for [Qwen Image](https://huggingface.co/docs/diffusers/main/en/api/pipelines/qwen) - [`LoraBaseMixin`] provides a base class with several utility methods to fuse, unfuse, unload, LoRAs and more. - - -To learn more about how to load LoRA weights, see the [LoRA](../../using-diffusers/loading_adapters#lora) loading guide. - - +> [!TIP] +> To learn more about how to load LoRA weights, see the [LoRA](../../using-diffusers/loading_adapters#lora) loading guide. ## LoraBaseMixin diff --git a/docs/source/en/api/loaders/peft.md b/docs/source/en/api/loaders/peft.md index a371ab9c8e..5508509c88 100644 --- a/docs/source/en/api/loaders/peft.md +++ b/docs/source/en/api/loaders/peft.md @@ -14,11 +14,8 @@ specific language governing permissions and limitations under the License. Diffusers supports loading adapters such as [LoRA](../../using-diffusers/loading_adapters) with the [PEFT](https://huggingface.co/docs/peft/index) library with the [`~loaders.peft.PeftAdapterMixin`] class. This allows modeling classes in Diffusers like [`UNet2DConditionModel`], [`SD3Transformer2DModel`] to operate with an adapter. - - -Refer to the [Inference with PEFT](../../tutorials/using_peft_for_inference.md) tutorial for an overview of how to use PEFT in Diffusers for inference. - - +> [!TIP] +> Refer to the [Inference with PEFT](../../tutorials/using_peft_for_inference.md) tutorial for an overview of how to use PEFT in Diffusers for inference. ## PeftAdapterMixin diff --git a/docs/source/en/api/loaders/textual_inversion.md b/docs/source/en/api/loaders/textual_inversion.md index 30d8f5b8d5..2cb54ce4ea 100644 --- a/docs/source/en/api/loaders/textual_inversion.md +++ b/docs/source/en/api/loaders/textual_inversion.md @@ -16,11 +16,8 @@ Textual Inversion is a training method for personalizing models by learning new [`TextualInversionLoaderMixin`] provides a function for loading Textual Inversion embeddings from Diffusers and Automatic1111 into the text encoder and loading a special token to activate the embeddings. - - -To learn more about how to load Textual Inversion embeddings, see the [Textual Inversion](../../using-diffusers/loading_adapters#textual-inversion) loading guide. - - +> [!TIP] +> To learn more about how to load Textual Inversion embeddings, see the [Textual Inversion](../../using-diffusers/loading_adapters#textual-inversion) loading guide. ## TextualInversionLoaderMixin diff --git a/docs/source/en/api/loaders/transformer_sd3.md b/docs/source/en/api/loaders/transformer_sd3.md index 0e7664cdd1..cc3ec0da14 100644 --- a/docs/source/en/api/loaders/transformer_sd3.md +++ b/docs/source/en/api/loaders/transformer_sd3.md @@ -16,11 +16,8 @@ This class is useful when *only* loading weights into a [`SD3Transformer2DModel` The [`SD3Transformer2DLoadersMixin`] class currently only loads IP-Adapter weights, but will be used in the future to save weights and load LoRAs. - - -To learn more about how to load LoRA weights, see the [LoRA](../../using-diffusers/loading_adapters#lora) loading guide. - - +> [!TIP] +> To learn more about how to load LoRA weights, see the [LoRA](../../using-diffusers/loading_adapters#lora) loading guide. ## SD3Transformer2DLoadersMixin diff --git a/docs/source/en/api/loaders/unet.md b/docs/source/en/api/loaders/unet.md index 51b4c4ef48..7450e03e58 100644 --- a/docs/source/en/api/loaders/unet.md +++ b/docs/source/en/api/loaders/unet.md @@ -16,11 +16,8 @@ Some training methods - like LoRA and Custom Diffusion - typically target the UN The [`UNet2DConditionLoadersMixin`] class provides functions for loading and saving weights, fusing and unfusing LoRAs, disabling and enabling LoRAs, and setting and deleting adapters. - - -To learn more about how to load LoRA weights, see the [LoRA](../../using-diffusers/loading_adapters#lora) loading guide. - - +> [!TIP] +> To learn more about how to load LoRA weights, see the [LoRA](../../using-diffusers/loading_adapters#lora) loading guide. ## UNet2DConditionLoadersMixin diff --git a/docs/source/en/api/models/consistency_decoder_vae.md b/docs/source/en/api/models/consistency_decoder_vae.md index cf4955a074..fe039df7f9 100644 --- a/docs/source/en/api/models/consistency_decoder_vae.md +++ b/docs/source/en/api/models/consistency_decoder_vae.md @@ -16,11 +16,8 @@ Consistency decoder can be used to decode the latents from the denoising UNet in The original codebase can be found at [openai/consistencydecoder](https://github.com/openai/consistencydecoder). - - -Inference is only supported for 2 iterations as of now. - - +> [!WARNING] +> Inference is only supported for 2 iterations as of now. The pipeline could not have been contributed without the help of [madebyollin](https://github.com/madebyollin) and [mrsteyk](https://github.com/mrsteyk) from [this issue](https://github.com/openai/consistencydecoder/issues/1). diff --git a/docs/source/en/api/models/transformer2d.md b/docs/source/en/api/models/transformer2d.md index 16ae6ace97..d8e0a858b0 100644 --- a/docs/source/en/api/models/transformer2d.md +++ b/docs/source/en/api/models/transformer2d.md @@ -22,11 +22,8 @@ When the input is **continuous**: When the input is **discrete**: - - -It is assumed one of the input classes is the masked latent pixel. The predicted classes of the unnoised image don't contain a prediction for the masked pixel because the unnoised image cannot be masked. - - +> [!TIP] +> It is assumed one of the input classes is the masked latent pixel. The predicted classes of the unnoised image don't contain a prediction for the masked pixel because the unnoised image cannot be masked. 1. Convert input (classes of latent pixels) to embeddings and apply positional embeddings. 2. Apply the Transformer blocks in the standard way. diff --git a/docs/source/en/api/outputs.md b/docs/source/en/api/outputs.md index a13bc89f2b..0fba1ab2fa 100644 --- a/docs/source/en/api/outputs.md +++ b/docs/source/en/api/outputs.md @@ -39,11 +39,8 @@ For instance, retrieving an image by indexing into it returns the tuple `(output outputs[:1] ``` - - -To check a specific pipeline or model output, refer to its corresponding API documentation. - - +> [!TIP] +> To check a specific pipeline or model output, refer to its corresponding API documentation. ## BaseOutput diff --git a/docs/source/en/api/pipelines/allegro.md b/docs/source/en/api/pipelines/allegro.md index 09313c2db0..a981fb1f94 100644 --- a/docs/source/en/api/pipelines/allegro.md +++ b/docs/source/en/api/pipelines/allegro.md @@ -17,11 +17,8 @@ The abstract from the paper is: *Significant advancements have been made in the field of video generation, with the open-source community contributing a wealth of research papers and tools for training high-quality models. However, despite these efforts, the available information and resources remain insufficient for achieving commercial-level performance. In this report, we open the black box and introduce Allegro, an advanced video generation model that excels in both quality and temporal consistency. We also highlight the current limitations in the field and present a comprehensive methodology for training high-performance, commercial-level video generation models, addressing key aspects such as data, model architecture, training pipeline, and evaluation. Our user study shows that Allegro surpasses existing open-source models and most commercial models, ranking just behind Hailuo and Kling. Code: https://github.com/rhymes-ai/Allegro , Model: https://huggingface.co/rhymes-ai/Allegro , Gallery: https://rhymes.ai/allegro_gallery .* - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## Quantization diff --git a/docs/source/en/api/pipelines/animatediff.md b/docs/source/en/api/pipelines/animatediff.md index aeec3254ca..f0188f3c36 100644 --- a/docs/source/en/api/pipelines/animatediff.md +++ b/docs/source/en/api/pipelines/animatediff.md @@ -102,11 +102,8 @@ Here are some sample outputs: - - -AnimateDiff tends to work better with finetuned Stable Diffusion models. If you plan on using a scheduler that can clip samples, make sure to disable it by setting `clip_sample=False` in the scheduler as this can also have an adverse effect on generated samples. Additionally, the AnimateDiff checkpoints can be sensitive to the beta schedule of the scheduler. We recommend setting this to `linear`. - - +> [!TIP] +> AnimateDiff tends to work better with finetuned Stable Diffusion models. If you plan on using a scheduler that can clip samples, make sure to disable it by setting `clip_sample=False` in the scheduler as this can also have an adverse effect on generated samples. Additionally, the AnimateDiff checkpoints can be sensitive to the beta schedule of the scheduler. We recommend setting this to `linear`. ### AnimateDiffControlNetPipeline @@ -799,17 +796,11 @@ frames = output.frames[0] export_to_gif(frames, "animation.gif") ``` - +> [!WARNING] +> FreeInit is not really free - the improved quality comes at the cost of extra computation. It requires sampling a few extra times depending on the `num_iters` parameter that is set when enabling it. Setting the `use_fast_sampling` parameter to `True` can improve the overall performance (at the cost of lower quality compared to when `use_fast_sampling=False` but still better results than vanilla video generation models). -FreeInit is not really free - the improved quality comes at the cost of extra computation. It requires sampling a few extra times depending on the `num_iters` parameter that is set when enabling it. Setting the `use_fast_sampling` parameter to `True` can improve the overall performance (at the cost of lower quality compared to when `use_fast_sampling=False` but still better results than vanilla video generation models). - - - - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. diff --git a/docs/source/en/api/pipelines/attend_and_excite.md b/docs/source/en/api/pipelines/attend_and_excite.md index b5ce3bb767..e7d1e1d2b8 100644 --- a/docs/source/en/api/pipelines/attend_and_excite.md +++ b/docs/source/en/api/pipelines/attend_and_excite.md @@ -23,11 +23,8 @@ The abstract from the paper is: You can find additional information about Attend-and-Excite on the [project page](https://attendandexcite.github.io/Attend-and-Excite/), the [original codebase](https://github.com/AttendAndExcite/Attend-and-Excite), or try it out in a [demo](https://huggingface.co/spaces/AttendAndExcite/Attend-and-Excite). - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## StableDiffusionAttendAndExcitePipeline diff --git a/docs/source/en/api/pipelines/audioldm.md b/docs/source/en/api/pipelines/audioldm.md index 6b143d2990..c8073a14ef 100644 --- a/docs/source/en/api/pipelines/audioldm.md +++ b/docs/source/en/api/pipelines/audioldm.md @@ -38,11 +38,8 @@ During inference: * The _quality_ of the predicted audio sample can be controlled by the `num_inference_steps` argument; higher steps give higher quality audio at the expense of slower inference. * The _length_ of the predicted audio sample can be controlled by varying the `audio_length_in_s` argument. - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## AudioLDMPipeline [[autodoc]] AudioLDMPipeline diff --git a/docs/source/en/api/pipelines/audioldm2.md b/docs/source/en/api/pipelines/audioldm2.md index 1a196099d7..45a9002ea0 100644 --- a/docs/source/en/api/pipelines/audioldm2.md +++ b/docs/source/en/api/pipelines/audioldm2.md @@ -58,11 +58,8 @@ See table below for details on the three checkpoints: The following example demonstrates how to construct good music and speech generation using the aforementioned tips: [example](https://huggingface.co/docs/diffusers/main/en/api/pipelines/audioldm2#diffusers.AudioLDM2Pipeline.__call__.example). - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## AudioLDM2Pipeline [[autodoc]] AudioLDM2Pipeline diff --git a/docs/source/en/api/pipelines/aura_flow.md b/docs/source/en/api/pipelines/aura_flow.md index 1d6002335c..67951859b9 100644 --- a/docs/source/en/api/pipelines/aura_flow.md +++ b/docs/source/en/api/pipelines/aura_flow.md @@ -16,11 +16,8 @@ AuraFlow is inspired by [Stable Diffusion 3](../pipelines/stable_diffusion/stabl It was developed by the Fal team and more details about it can be found in [this blog post](https://blog.fal.ai/auraflow/). - - -AuraFlow can be quite expensive to run on consumer hardware devices. However, you can perform a suite of optimizations to run it faster and in a more memory-friendly manner. Check out [this section](https://huggingface.co/blog/sd3#memory-optimizations-for-sd3) for more details. - - +> [!TIP] +> AuraFlow can be quite expensive to run on consumer hardware devices. However, you can perform a suite of optimizations to run it faster and in a more memory-friendly manner. Check out [this section](https://huggingface.co/blog/sd3#memory-optimizations-for-sd3) for more details. ## Quantization diff --git a/docs/source/en/api/pipelines/blip_diffusion.md b/docs/source/en/api/pipelines/blip_diffusion.md index d94281a4a9..b9c6ed7b5f 100644 --- a/docs/source/en/api/pipelines/blip_diffusion.md +++ b/docs/source/en/api/pipelines/blip_diffusion.md @@ -26,11 +26,8 @@ The original codebase can be found at [salesforce/LAVIS](https://github.com/sale `BlipDiffusionPipeline` and `BlipDiffusionControlNetPipeline` were contributed by [`ayushtues`](https://github.com/ayushtues/). - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## BlipDiffusionPipeline diff --git a/docs/source/en/api/pipelines/chroma.md b/docs/source/en/api/pipelines/chroma.md index 40e290e4bd..df03fbb325 100644 --- a/docs/source/en/api/pipelines/chroma.md +++ b/docs/source/en/api/pipelines/chroma.md @@ -21,11 +21,8 @@ Chroma is a text to image generation model based on Flux. Original model checkpoints for Chroma can be found [here](https://huggingface.co/lodestones/Chroma). - - -Chroma can use all the same optimizations as Flux. - - +> [!TIP] +> Chroma can use all the same optimizations as Flux. ## Inference diff --git a/docs/source/en/api/pipelines/cogview3.md b/docs/source/en/api/pipelines/cogview3.md index 0180fee300..5ee02e1a70 100644 --- a/docs/source/en/api/pipelines/cogview3.md +++ b/docs/source/en/api/pipelines/cogview3.md @@ -21,11 +21,8 @@ The abstract from the paper is: *Recent advancements in text-to-image generative systems have been largely driven by diffusion models. However, single-stage text-to-image diffusion models still face challenges, in terms of computational efficiency and the refinement of image details. To tackle the issue, we propose CogView3, an innovative cascaded framework that enhances the performance of text-to-image diffusion. CogView3 is the first model implementing relay diffusion in the realm of text-to-image generation, executing the task by first creating low-resolution images and subsequently applying relay-based super-resolution. This methodology not only results in competitive text-to-image outputs but also greatly reduces both training and inference costs. Our experimental results demonstrate that CogView3 outperforms SDXL, the current state-of-the-art open-source text-to-image diffusion model, by 77.0% in human evaluations, all while requiring only about 1/2 of the inference time. The distilled variant of CogView3 achieves comparable performance while only utilizing 1/10 of the inference time by SDXL.* - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. This pipeline was contributed by [zRzRzRzRzRzRzR](https://github.com/zRzRzRzRzRzRzR). The original codebase can be found [here](https://huggingface.co/THUDM). The original weights can be found under [hf.co/THUDM](https://huggingface.co/THUDM). diff --git a/docs/source/en/api/pipelines/cogview4.md b/docs/source/en/api/pipelines/cogview4.md index 50ba5baa62..7857dc8c94 100644 --- a/docs/source/en/api/pipelines/cogview4.md +++ b/docs/source/en/api/pipelines/cogview4.md @@ -15,11 +15,8 @@ # CogView4 - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. This pipeline was contributed by [zRzRzRzRzRzRzR](https://github.com/zRzRzRzRzRzRzR). The original codebase can be found [here](https://huggingface.co/THUDM). The original weights can be found under [hf.co/THUDM](https://huggingface.co/THUDM). diff --git a/docs/source/en/api/pipelines/consisid.md b/docs/source/en/api/pipelines/consisid.md index db6b5e59ac..bba0472924 100644 --- a/docs/source/en/api/pipelines/consisid.md +++ b/docs/source/en/api/pipelines/consisid.md @@ -25,11 +25,8 @@ The abstract from the paper is: *Identity-preserving text-to-video (IPT2V) generation aims to create high-fidelity videos with consistent human identity. It is an important task in video generation but remains an open problem for generative models. This paper pushes the technical frontier of IPT2V in two directions that have not been resolved in the literature: (1) A tuning-free pipeline without tedious case-by-case finetuning, and (2) A frequency-aware heuristic identity-preserving Diffusion Transformer (DiT)-based control scheme. To achieve these goals, we propose **ConsisID**, a tuning-free DiT-based controllable IPT2V model to keep human-**id**entity **consis**tent in the generated video. Inspired by prior findings in frequency analysis of vision/diffusion transformers, it employs identity-control signals in the frequency domain, where facial features can be decomposed into low-frequency global features (e.g., profile, proportions) and high-frequency intrinsic features (e.g., identity markers that remain unaffected by pose changes). First, from a low-frequency perspective, we introduce a global facial extractor, which encodes the reference image and facial key points into a latent space, generating features enriched with low-frequency information. These features are then integrated into the shallow layers of the network to alleviate training challenges associated with DiT. Second, from a high-frequency perspective, we design a local facial extractor to capture high-frequency details and inject them into the transformer blocks, enhancing the model's ability to preserve fine-grained features. To leverage the frequency information for identity preservation, we propose a hierarchical training strategy, transforming a vanilla pre-trained video generation model into an IPT2V model. Extensive experiments demonstrate that our frequency-aware heuristic scheme provides an optimal control solution for DiT-based models. Thanks to this scheme, our **ConsisID** achieves excellent results in generating high-quality, identity-preserving videos, making strides towards more effective IPT2V. The model weight of ConsID is publicly available at https://github.com/PKU-YuanGroup/ConsisID.* - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers.md) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading.md#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers.md) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading.md#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. This pipeline was contributed by [SHYuanBest](https://github.com/SHYuanBest). The original codebase can be found [here](https://github.com/PKU-YuanGroup/ConsisID). The original weights can be found under [hf.co/BestWishYsh](https://huggingface.co/BestWishYsh). diff --git a/docs/source/en/api/pipelines/control_flux_inpaint.md b/docs/source/en/api/pipelines/control_flux_inpaint.md index 03a4fbebb8..4b087f20ef 100644 --- a/docs/source/en/api/pipelines/control_flux_inpaint.md +++ b/docs/source/en/api/pipelines/control_flux_inpaint.md @@ -26,11 +26,8 @@ FLUX.1 Depth and Canny [dev] is a 12 billion parameter rectified flow transforme | Canny | [Black Forest Labs](https://huggingface.co/black-forest-labs) | [Link](https://huggingface.co/black-forest-labs/FLUX.1-Canny-dev) | - - -Flux can be quite expensive to run on consumer hardware devices. However, you can perform a suite of optimizations to run it faster and in a more memory-friendly manner. Check out [this section](https://huggingface.co/blog/sd3#memory-optimizations-for-sd3) for more details. Additionally, Flux can benefit from quantization for memory efficiency with a trade-off in inference latency. Refer to [this blog post](https://huggingface.co/blog/quanto-diffusers) to learn more. For an exhaustive list of resources, check out [this gist](https://gist.github.com/sayakpaul/b664605caf0aa3bf8585ab109dd5ac9c). - - +> [!TIP] +> Flux can be quite expensive to run on consumer hardware devices. However, you can perform a suite of optimizations to run it faster and in a more memory-friendly manner. Check out [this section](https://huggingface.co/blog/sd3#memory-optimizations-for-sd3) for more details. Additionally, Flux can benefit from quantization for memory efficiency with a trade-off in inference latency. Refer to [this blog post](https://huggingface.co/blog/quanto-diffusers) to learn more. For an exhaustive list of resources, check out [this gist](https://gist.github.com/sayakpaul/b664605caf0aa3bf8585ab109dd5ac9c). ```python import torch diff --git a/docs/source/en/api/pipelines/controlnet.md b/docs/source/en/api/pipelines/controlnet.md index 2a654a3735..afc0a4653e 100644 --- a/docs/source/en/api/pipelines/controlnet.md +++ b/docs/source/en/api/pipelines/controlnet.md @@ -28,11 +28,8 @@ This model was contributed by [takuma104](https://huggingface.co/takuma104). ❤ The original codebase can be found at [lllyasviel/ControlNet](https://github.com/lllyasviel/ControlNet), and you can find official ControlNet checkpoints on [lllyasviel's](https://huggingface.co/lllyasviel) Hub profile. - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## StableDiffusionControlNetPipeline [[autodoc]] StableDiffusionControlNetPipeline diff --git a/docs/source/en/api/pipelines/controlnet_flux.md b/docs/source/en/api/pipelines/controlnet_flux.md index 9feb736523..ff38ca3f2c 100644 --- a/docs/source/en/api/pipelines/controlnet_flux.md +++ b/docs/source/en/api/pipelines/controlnet_flux.md @@ -44,11 +44,8 @@ XLabs ControlNets are also supported, which was contributed by the [XLabs team]( | HED | [The XLabs Team](https://huggingface.co/XLabs-AI) | [Link](https://huggingface.co/XLabs-AI/flux-controlnet-hed-diffusers) | - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## FluxControlNetPipeline [[autodoc]] FluxControlNetPipeline diff --git a/docs/source/en/api/pipelines/controlnet_hunyuandit.md b/docs/source/en/api/pipelines/controlnet_hunyuandit.md index c79b2dbf65..88dc2de10a 100644 --- a/docs/source/en/api/pipelines/controlnet_hunyuandit.md +++ b/docs/source/en/api/pipelines/controlnet_hunyuandit.md @@ -24,11 +24,8 @@ The abstract from the paper is: This code is implemented by Tencent Hunyuan Team. You can find pre-trained checkpoints for Hunyuan-DiT ControlNets on [Tencent Hunyuan](https://huggingface.co/Tencent-Hunyuan). - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## HunyuanDiTControlNetPipeline [[autodoc]] HunyuanDiTControlNetPipeline diff --git a/docs/source/en/api/pipelines/controlnet_sd3.md b/docs/source/en/api/pipelines/controlnet_sd3.md index 067c1c6b01..8cdada9edf 100644 --- a/docs/source/en/api/pipelines/controlnet_sd3.md +++ b/docs/source/en/api/pipelines/controlnet_sd3.md @@ -38,11 +38,8 @@ This controlnet code is mainly implemented by [The InstantX Team](https://huggin | Inpainting | [The AlimamaCreative Team](https://huggingface.co/alimama-creative) | [link](https://huggingface.co/alimama-creative/SD3-Controlnet-Inpainting) | - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## StableDiffusion3ControlNetPipeline [[autodoc]] StableDiffusion3ControlNetPipeline diff --git a/docs/source/en/api/pipelines/controlnet_sdxl.md b/docs/source/en/api/pipelines/controlnet_sdxl.md index cb0554a1cc..89fc1c3897 100644 --- a/docs/source/en/api/pipelines/controlnet_sdxl.md +++ b/docs/source/en/api/pipelines/controlnet_sdxl.md @@ -26,19 +26,13 @@ The abstract from the paper is: You can find additional smaller Stable Diffusion XL (SDXL) ControlNet checkpoints from the 🤗 [Diffusers](https://huggingface.co/diffusers) Hub organization, and browse [community-trained](https://huggingface.co/models?other=stable-diffusion-xl&other=controlnet) checkpoints on the Hub. - - -🧪 Many of the SDXL ControlNet checkpoints are experimental, and there is a lot of room for improvement. Feel free to open an [Issue](https://github.com/huggingface/diffusers/issues/new/choose) and leave us feedback on how we can improve! - - +> [!WARNING] +> 🧪 Many of the SDXL ControlNet checkpoints are experimental, and there is a lot of room for improvement. Feel free to open an [Issue](https://github.com/huggingface/diffusers/issues/new/choose) and leave us feedback on how we can improve! If you don't see a checkpoint you're interested in, you can train your own SDXL ControlNet with our [training script](../../../../../examples/controlnet/README_sdxl). - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## StableDiffusionXLControlNetPipeline [[autodoc]] StableDiffusionXLControlNetPipeline diff --git a/docs/source/en/api/pipelines/controlnetxs.md b/docs/source/en/api/pipelines/controlnetxs.md index aea8cb2e86..d44fb0cf0f 100644 --- a/docs/source/en/api/pipelines/controlnetxs.md +++ b/docs/source/en/api/pipelines/controlnetxs.md @@ -31,11 +31,8 @@ Here's the overview from the [project page](https://vislearn.github.io/ControlNe This model was contributed by [UmerHA](https://twitter.com/UmerHAdil). ❤️ - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## StableDiffusionControlNetXSPipeline [[autodoc]] StableDiffusionControlNetXSPipeline diff --git a/docs/source/en/api/pipelines/controlnetxs_sdxl.md b/docs/source/en/api/pipelines/controlnetxs_sdxl.md index 76937b16c5..7ae0e2a2a1 100644 --- a/docs/source/en/api/pipelines/controlnetxs_sdxl.md +++ b/docs/source/en/api/pipelines/controlnetxs_sdxl.md @@ -27,17 +27,11 @@ Here's the overview from the [project page](https://vislearn.github.io/ControlNe This model was contributed by [UmerHA](https://twitter.com/UmerHAdil). ❤️ - +> [!WARNING] +> 🧪 Many of the SDXL ControlNet checkpoints are experimental, and there is a lot of room for improvement. Feel free to open an [Issue](https://github.com/huggingface/diffusers/issues/new/choose) and leave us feedback on how we can improve! -🧪 Many of the SDXL ControlNet checkpoints are experimental, and there is a lot of room for improvement. Feel free to open an [Issue](https://github.com/huggingface/diffusers/issues/new/choose) and leave us feedback on how we can improve! - - - - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## StableDiffusionXLControlNetXSPipeline [[autodoc]] StableDiffusionXLControlNetXSPipeline diff --git a/docs/source/en/api/pipelines/cosmos.md b/docs/source/en/api/pipelines/cosmos.md index dba807c5ce..fb9453480e 100644 --- a/docs/source/en/api/pipelines/cosmos.md +++ b/docs/source/en/api/pipelines/cosmos.md @@ -18,11 +18,8 @@ *Physical AI needs to be trained digitally first. It needs a digital twin of itself, the policy model, and a digital twin of the world, the world model. In this paper, we present the Cosmos World Foundation Model Platform to help developers build customized world models for their Physical AI setups. We position a world foundation model as a general-purpose world model that can be fine-tuned into customized world models for downstream applications. Our platform covers a video curation pipeline, pre-trained world foundation models, examples of post-training of pre-trained world foundation models, and video tokenizers. To help Physical AI builders solve the most critical problems of our society, we make our platform open-source and our models open-weight with permissive licenses available via https://github.com/NVIDIA/Cosmos.* - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## Loading original format checkpoints diff --git a/docs/source/en/api/pipelines/dance_diffusion.md b/docs/source/en/api/pipelines/dance_diffusion.md index 5805561e49..0434f63195 100644 --- a/docs/source/en/api/pipelines/dance_diffusion.md +++ b/docs/source/en/api/pipelines/dance_diffusion.md @@ -20,11 +20,8 @@ specific language governing permissions and limitations under the License. Dance Diffusion is the first in a suite of generative audio tools for producers and musicians released by [Harmonai](https://github.com/Harmonai-org). - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## DanceDiffusionPipeline [[autodoc]] DanceDiffusionPipeline diff --git a/docs/source/en/api/pipelines/ddpm.md b/docs/source/en/api/pipelines/ddpm.md index 716cf73275..63c2fcaf89 100644 --- a/docs/source/en/api/pipelines/ddpm.md +++ b/docs/source/en/api/pipelines/ddpm.md @@ -20,11 +20,8 @@ The abstract from the paper is: The original codebase can be found at [hohonathanho/diffusion](https://github.com/hojonathanho/diffusion). - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. # DDPMPipeline [[autodoc]] DDPMPipeline diff --git a/docs/source/en/api/pipelines/dit.md b/docs/source/en/api/pipelines/dit.md index e87058899b..16d0c99961 100644 --- a/docs/source/en/api/pipelines/dit.md +++ b/docs/source/en/api/pipelines/dit.md @@ -20,11 +20,8 @@ The abstract from the paper is: The original codebase can be found at [facebookresearch/dit](https://github.com/facebookresearch/dit). - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## DiTPipeline [[autodoc]] DiTPipeline diff --git a/docs/source/en/api/pipelines/flux.md b/docs/source/en/api/pipelines/flux.md index bb72758222..1a89de98e4 100644 --- a/docs/source/en/api/pipelines/flux.md +++ b/docs/source/en/api/pipelines/flux.md @@ -21,13 +21,10 @@ Flux is a series of text-to-image generation models based on diffusion transform Original model checkpoints for Flux can be found [here](https://huggingface.co/black-forest-labs). Original inference code can be found [here](https://github.com/black-forest-labs/flux). - - -Flux can be quite expensive to run on consumer hardware devices. However, you can perform a suite of optimizations to run it faster and in a more memory-friendly manner. Check out [this section](https://huggingface.co/blog/sd3#memory-optimizations-for-sd3) for more details. Additionally, Flux can benefit from quantization for memory efficiency with a trade-off in inference latency. Refer to [this blog post](https://huggingface.co/blog/quanto-diffusers) to learn more. For an exhaustive list of resources, check out [this gist](https://gist.github.com/sayakpaul/b664605caf0aa3bf8585ab109dd5ac9c). - -[Caching](../../optimization/cache) may also speed up inference by storing and reusing intermediate outputs. - - +> [!TIP] +> Flux can be quite expensive to run on consumer hardware devices. However, you can perform a suite of optimizations to run it faster and in a more memory-friendly manner. Check out [this section](https://huggingface.co/blog/sd3#memory-optimizations-for-sd3) for more details. Additionally, Flux can benefit from quantization for memory efficiency with a trade-off in inference latency. Refer to [this blog post](https://huggingface.co/blog/quanto-diffusers) to learn more. For an exhaustive list of resources, check out [this gist](https://gist.github.com/sayakpaul/b664605caf0aa3bf8585ab109dd5ac9c). +> +> [Caching](../../optimization/cache) may also speed up inference by storing and reusing intermediate outputs. Flux comes in the following variants: @@ -420,11 +417,8 @@ When unloading the Control LoRA weights, call `pipe.unload_lora_weights(reset_to ## IP-Adapter - - -Check out [IP-Adapter](../../../using-diffusers/ip_adapter) to learn more about how IP-Adapters work. - - +> [!TIP] +> Check out [IP-Adapter](../../../using-diffusers/ip_adapter) to learn more about how IP-Adapters work. An IP-Adapter lets you prompt Flux with images, in addition to the text prompt. This is especially useful when describing complex concepts that are difficult to articulate through text alone and you have reference images. @@ -604,9 +598,8 @@ image.save("flux.png") The `FluxTransformer2DModel` supports loading checkpoints in the original format shipped by Black Forest Labs. This is also useful when trying to load finetunes or quantized versions of the models that have been published by the community. - -`FP8` inference can be brittle depending on the GPU type, CUDA version, and `torch` version that you are using. It is recommended that you use the `optimum-quanto` library in order to run FP8 inference on your machine. - +> [!TIP] +> `FP8` inference can be brittle depending on the GPU type, CUDA version, and `torch` version that you are using. It is recommended that you use the `optimum-quanto` library in order to run FP8 inference on your machine. The following example demonstrates how to run Flux with less than 16GB of VRAM. diff --git a/docs/source/en/api/pipelines/framepack.md b/docs/source/en/api/pipelines/framepack.md index ba7b2d0dc0..a25cfe24a4 100644 --- a/docs/source/en/api/pipelines/framepack.md +++ b/docs/source/en/api/pipelines/framepack.md @@ -22,11 +22,8 @@ *We present a neural network structure, FramePack, to train next-frame (or next-frame-section) prediction models for video generation. The FramePack compresses input frames to make the transformer context length a fixed number regardless of the video length. As a result, we are able to process a large number of frames using video diffusion with computation bottleneck similar to image diffusion. This also makes the training video batch sizes significantly higher (batch sizes become comparable to image diffusion training). We also propose an anti-drifting sampling method that generates frames in inverted temporal order with early-established endpoints to avoid exposure bias (error accumulation over iterations). Finally, we show that existing video diffusion models can be finetuned with FramePack, and their visual quality may be improved because the next-frame prediction supports more balanced diffusion schedulers with less extreme flow shift timesteps.* - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## Available models diff --git a/docs/source/en/api/pipelines/hidream.md b/docs/source/en/api/pipelines/hidream.md index 9848612c33..acfcef93e0 100644 --- a/docs/source/en/api/pipelines/hidream.md +++ b/docs/source/en/api/pipelines/hidream.md @@ -16,11 +16,8 @@ [HiDream-I1](https://huggingface.co/HiDream-ai) by HiDream.ai - - -[Caching](../../optimization/cache) may also speed up inference by storing and reusing intermediate outputs. - - +> [!TIP] +> [Caching](../../optimization/cache) may also speed up inference by storing and reusing intermediate outputs. ## Available models diff --git a/docs/source/en/api/pipelines/hunyuandit.md b/docs/source/en/api/pipelines/hunyuandit.md index 07e869ba95..3f4db66c6c 100644 --- a/docs/source/en/api/pipelines/hunyuandit.md +++ b/docs/source/en/api/pipelines/hunyuandit.md @@ -28,17 +28,11 @@ HunyuanDiT has the following components: * It uses a diffusion transformer as the backbone * It combines two text encoders, a bilingual CLIP and a multilingual T5 encoder - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - - - - -You can further improve generation quality by passing the generated image from [`HungyuanDiTPipeline`] to the [SDXL refiner](../../using-diffusers/sdxl#base-to-refiner-model) model. - - +> [!TIP] +> You can further improve generation quality by passing the generated image from [`HungyuanDiTPipeline`] to the [SDXL refiner](../../using-diffusers/sdxl#base-to-refiner-model) model. ## Optimization diff --git a/docs/source/en/api/pipelines/i2vgenxl.md b/docs/source/en/api/pipelines/i2vgenxl.md index 76a51a6cd5..711a5625f9 100644 --- a/docs/source/en/api/pipelines/i2vgenxl.md +++ b/docs/source/en/api/pipelines/i2vgenxl.md @@ -23,11 +23,8 @@ The abstract from the paper is: The original codebase can be found [here](https://github.com/ali-vilab/i2vgen-xl/). The model checkpoints can be found [here](https://huggingface.co/ali-vilab/). - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. Also, to know more about reducing the memory usage of this pipeline, refer to the ["Reduce memory usage"] section [here](../../using-diffusers/svd#reduce-memory-usage). - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. Also, to know more about reducing the memory usage of this pipeline, refer to the ["Reduce memory usage"] section [here](../../using-diffusers/svd#reduce-memory-usage). Sample output with I2VGenXL: diff --git a/docs/source/en/api/pipelines/kandinsky.md b/docs/source/en/api/pipelines/kandinsky.md index 90c76954ab..7717f2db69 100644 --- a/docs/source/en/api/pipelines/kandinsky.md +++ b/docs/source/en/api/pipelines/kandinsky.md @@ -17,17 +17,11 @@ The description from it's GitHub page is: The original codebase can be found at [ai-forever/Kandinsky-2](https://github.com/ai-forever/Kandinsky-2). - +> [!TIP] +> Check out the [Kandinsky Community](https://huggingface.co/kandinsky-community) organization on the Hub for the official model checkpoints for tasks like text-to-image, image-to-image, and inpainting. -Check out the [Kandinsky Community](https://huggingface.co/kandinsky-community) organization on the Hub for the official model checkpoints for tasks like text-to-image, image-to-image, and inpainting. - - - - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## KandinskyPriorPipeline diff --git a/docs/source/en/api/pipelines/kandinsky3.md b/docs/source/en/api/pipelines/kandinsky3.md index 1727387c4a..f08afa8879 100644 --- a/docs/source/en/api/pipelines/kandinsky3.md +++ b/docs/source/en/api/pipelines/kandinsky3.md @@ -28,17 +28,11 @@ Its architecture includes 3 main components: The original codebase can be found at [ai-forever/Kandinsky-3](https://github.com/ai-forever/Kandinsky-3). - +> [!TIP] +> Check out the [Kandinsky Community](https://huggingface.co/kandinsky-community) organization on the Hub for the official model checkpoints for tasks like text-to-image, image-to-image, and inpainting. -Check out the [Kandinsky Community](https://huggingface.co/kandinsky-community) organization on the Hub for the official model checkpoints for tasks like text-to-image, image-to-image, and inpainting. - - - - - -Make sure to check out the schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## Kandinsky3Pipeline diff --git a/docs/source/en/api/pipelines/kandinsky_v22.md b/docs/source/en/api/pipelines/kandinsky_v22.md index e68c094e23..0e0ed80db6 100644 --- a/docs/source/en/api/pipelines/kandinsky_v22.md +++ b/docs/source/en/api/pipelines/kandinsky_v22.md @@ -17,17 +17,11 @@ The description from it's GitHub page is: The original codebase can be found at [ai-forever/Kandinsky-2](https://github.com/ai-forever/Kandinsky-2). - +> [!TIP] +> Check out the [Kandinsky Community](https://huggingface.co/kandinsky-community) organization on the Hub for the official model checkpoints for tasks like text-to-image, image-to-image, and inpainting. -Check out the [Kandinsky Community](https://huggingface.co/kandinsky-community) organization on the Hub for the official model checkpoints for tasks like text-to-image, image-to-image, and inpainting. - - - - - -Make sure to check out the schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## KandinskyV22PriorPipeline diff --git a/docs/source/en/api/pipelines/kolors.md b/docs/source/en/api/pipelines/kolors.md index 048f6c1de9..b4c83fe134 100644 --- a/docs/source/en/api/pipelines/kolors.md +++ b/docs/source/en/api/pipelines/kolors.md @@ -50,17 +50,11 @@ image.save("kolors_sample.png") Kolors needs a different IP Adapter to work, and it uses [Openai-CLIP-336](https://huggingface.co/openai/clip-vit-large-patch14-336) as an image encoder. - +> [!TIP] +> Using an IP Adapter with Kolors requires more than 24GB of VRAM. To use it, we recommend using [`~DiffusionPipeline.enable_model_cpu_offload`] on consumer GPUs. -Using an IP Adapter with Kolors requires more than 24GB of VRAM. To use it, we recommend using [`~DiffusionPipeline.enable_model_cpu_offload`] on consumer GPUs. - - - - - -While Kolors is integrated in Diffusers, you need to load the image encoder from a revision to use the safetensor files. You can still use the main branch of the original repository if you're comfortable loading pickle checkpoints. - - +> [!TIP] +> While Kolors is integrated in Diffusers, you need to load the image encoder from a revision to use the safetensor files. You can still use the main branch of the original repository if you're comfortable loading pickle checkpoints. ```python import torch diff --git a/docs/source/en/api/pipelines/latent_diffusion.md b/docs/source/en/api/pipelines/latent_diffusion.md index 5489d673f5..cefed90e86 100644 --- a/docs/source/en/api/pipelines/latent_diffusion.md +++ b/docs/source/en/api/pipelines/latent_diffusion.md @@ -20,11 +20,8 @@ The abstract from the paper is: The original codebase can be found at [CompVis/latent-diffusion](https://github.com/CompVis/latent-diffusion). - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## LDMTextToImagePipeline [[autodoc]] LDMTextToImagePipeline diff --git a/docs/source/en/api/pipelines/latte.md b/docs/source/en/api/pipelines/latte.md index 9d4d12dd4e..c8438c668a 100644 --- a/docs/source/en/api/pipelines/latte.md +++ b/docs/source/en/api/pipelines/latte.md @@ -26,11 +26,8 @@ The abstract from the paper is: This pipeline was contributed by [maxin-cn](https://github.com/maxin-cn). The original codebase can be found [here](https://github.com/Vchitect/Latte). The original weights can be found under [hf.co/maxin-cn](https://huggingface.co/maxin-cn). - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ### Inference diff --git a/docs/source/en/api/pipelines/ledits_pp.md b/docs/source/en/api/pipelines/ledits_pp.md index 7c08971aa8..103bcf3798 100644 --- a/docs/source/en/api/pipelines/ledits_pp.md +++ b/docs/source/en/api/pipelines/ledits_pp.md @@ -22,16 +22,12 @@ The abstract from the paper is: *Text-to-image diffusion models have recently received increasing interest for their astonishing ability to produce high-fidelity images from solely text inputs. Subsequent research efforts aim to exploit and apply their capabilities to real image editing. However, existing image-to-image methods are often inefficient, imprecise, and of limited versatility. They either require time-consuming fine-tuning, deviate unnecessarily strongly from the input image, and/or lack support for multiple, simultaneous edits. To address these issues, we introduce LEDITS++, an efficient yet versatile and precise textual image manipulation technique. LEDITS++'s novel inversion approach requires no tuning nor optimization and produces high-fidelity results with a few diffusion steps. Second, our methodology supports multiple simultaneous edits and is architecture-agnostic. Third, we use a novel implicit masking technique that limits changes to relevant image regions. We propose the novel TEdBench++ benchmark as part of our exhaustive evaluation. Our results demonstrate the capabilities of LEDITS++ and its improvements over previous methods. The project page is available at https://leditsplusplus-project.static.hf.space .* - +> [!TIP] +> You can find additional information about LEDITS++ on the [project page](https://leditsplusplus-project.static.hf.space/index.html) and try it out in a [demo](https://huggingface.co/spaces/editing-images/leditsplusplus). -You can find additional information about LEDITS++ on the [project page](https://leditsplusplus-project.static.hf.space/index.html) and try it out in a [demo](https://huggingface.co/spaces/editing-images/leditsplusplus). - - - - -Due to some backward compatibility issues with the current diffusers implementation of [`~schedulers.DPMSolverMultistepScheduler`] this implementation of LEdits++ can no longer guarantee perfect inversion. -This issue is unlikely to have any noticeable effects on applied use-cases. However, we provide an alternative implementation that guarantees perfect inversion in a dedicated [GitHub repo](https://github.com/ml-research/ledits_pp). - +> [!WARNING] +> Due to some backward compatibility issues with the current diffusers implementation of [`~schedulers.DPMSolverMultistepScheduler`] this implementation of LEdits++ can no longer guarantee perfect inversion. +> This issue is unlikely to have any noticeable effects on applied use-cases. However, we provide an alternative implementation that guarantees perfect inversion in a dedicated [GitHub repo](https://github.com/ml-research/ledits_pp). We provide two distinct pipelines based on different pre-trained models. diff --git a/docs/source/en/api/pipelines/lumina.md b/docs/source/en/api/pipelines/lumina.md index 3bd3d9f8e0..0a236d213d 100644 --- a/docs/source/en/api/pipelines/lumina.md +++ b/docs/source/en/api/pipelines/lumina.md @@ -45,11 +45,8 @@ Lumina-T2X has the following components: This pipeline was contributed by [PommesPeter](https://github.com/PommesPeter). The original codebase can be found [here](https://github.com/Alpha-VLLM/Lumina-T2X). The original weights can be found under [hf.co/Alpha-VLLM](https://huggingface.co/Alpha-VLLM). - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ### Inference (Text-to-Image) diff --git a/docs/source/en/api/pipelines/lumina2.md b/docs/source/en/api/pipelines/lumina2.md index 092d7cde2e..0c4e793404 100644 --- a/docs/source/en/api/pipelines/lumina2.md +++ b/docs/source/en/api/pipelines/lumina2.md @@ -24,11 +24,8 @@ The abstract from the paper is: *We introduce Lumina-Image 2.0, an advanced text-to-image model that surpasses previous state-of-the-art methods across multiple benchmarks, while also shedding light on its potential to evolve into a generalist vision intelligence model. Lumina-Image 2.0 exhibits three key properties: (1) Unification – it adopts a unified architecture that treats text and image tokens as a joint sequence, enabling natural cross-modal interactions and facilitating task expansion. Besides, since high-quality captioners can provide semantically better-aligned text-image training pairs, we introduce a unified captioning system, UniCaptioner, which generates comprehensive and precise captions for the model. This not only accelerates model convergence but also enhances prompt adherence, variable-length prompt handling, and task generalization via prompt templates. (2) Efficiency – to improve the efficiency of the unified architecture, we develop a set of optimization techniques that improve semantic learning and fine-grained texture generation during training while incorporating inference-time acceleration strategies without compromising image quality. (3) Transparency – we open-source all training details, code, and models to ensure full reproducibility, aiming to bridge the gap between well-resourced closed-source research teams and independent developers.* - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## Using Single File loading with Lumina Image 2.0 diff --git a/docs/source/en/api/pipelines/marigold.md b/docs/source/en/api/pipelines/marigold.md index e9ca0df067..81e103afeb 100644 --- a/docs/source/en/api/pipelines/marigold.md +++ b/docs/source/en/api/pipelines/marigold.md @@ -45,14 +45,11 @@ This work expanded Marigold to support new modalities such as **Surface Normals* (IID), introduced a training protocol for **Latent Consistency Models** (LCM), and demonstrated **High-Resolution** (HR) processing capability. - - -The early Marigold models (`v1-0` and earlier) were optimized for best results with at least 10 inference steps. -LCM models were later developed to enable high-quality inference in just 1 to 4 steps. -Marigold models `v1-1` and later use the DDIM scheduler to achieve optimal -results in as few as 1 to 4 steps. - - +> [!TIP] +> The early Marigold models (`v1-0` and earlier) were optimized for best results with at least 10 inference steps. +> LCM models were later developed to enable high-quality inference in just 1 to 4 steps. +> Marigold models `v1-1` and later use the DDIM scheduler to achieve optimal +> results in as few as 1 to 4 steps. ## Available Pipelines @@ -80,27 +77,21 @@ The following is a summary of the recommended checkpoints, all of which produce | [prs-eth/marigold-iid-appearance-v1-1](https://huggingface.co/prs-eth/marigold-iid-appearance-v1-1) | Intrinsics | InteriorVerse decomposition is comprised of Albedo and two BRDF material properties: Roughness and Metallicity. | | [prs-eth/marigold-iid-lighting-v1-1](https://huggingface.co/prs-eth/marigold-iid-lighting-v1-1) | Intrinsics | HyperSim decomposition of an image  \\(I\\)  is comprised of Albedo  \\(A\\), Diffuse shading  \\(S\\), and Non-diffuse residual  \\(R\\):  \\(I = A*S+R\\). | - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff +> between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to +> efficiently load the same components into multiple pipelines. +> Also, to know more about reducing the memory usage of this pipeline, refer to the ["Reduce memory usage"] section +> [here](../../using-diffusers/svd#reduce-memory-usage). -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff -between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to -efficiently load the same components into multiple pipelines. -Also, to know more about reducing the memory usage of this pipeline, refer to the ["Reduce memory usage"] section -[here](../../using-diffusers/svd#reduce-memory-usage). - - - - - -Marigold pipelines were designed and tested with the scheduler embedded in the model checkpoint. -The optimal number of inference steps varies by scheduler, with no universal value that works best across all cases. -To accommodate this, the `num_inference_steps` parameter in the pipeline's `__call__` method defaults to `None` (see the -API reference). -Unless set explicitly, it inherits the value from the `default_denoising_steps` field in the checkpoint configuration -file (`model_index.json`). -This ensures high-quality predictions when invoking the pipeline with only the `image` argument. - - +> [!WARNING] +> Marigold pipelines were designed and tested with the scheduler embedded in the model checkpoint. +> The optimal number of inference steps varies by scheduler, with no universal value that works best across all cases. +> To accommodate this, the `num_inference_steps` parameter in the pipeline's `__call__` method defaults to `None` (see the +> API reference). +> Unless set explicitly, it inherits the value from the `default_denoising_steps` field in the checkpoint configuration +> file (`model_index.json`). +> This ensures high-quality predictions when invoking the pipeline with only the `image` argument. See also Marigold [usage examples](../../using-diffusers/marigold_usage). diff --git a/docs/source/en/api/pipelines/mochi.md b/docs/source/en/api/pipelines/mochi.md index f1260b07b0..f19a9bd575 100644 --- a/docs/source/en/api/pipelines/mochi.md +++ b/docs/source/en/api/pipelines/mochi.md @@ -121,15 +121,13 @@ export_to_video(frames, "mochi.mp4", fps=30) The [Genmo Mochi implementation](https://github.com/genmoai/mochi/tree/main) uses different precision values for each stage in the inference process. The text encoder and VAE use `torch.float32`, while the DiT uses `torch.bfloat16` with the [attention kernel](https://pytorch.org/docs/stable/generated/torch.nn.attention.sdpa_kernel.html#torch.nn.attention.sdpa_kernel) set to `EFFICIENT_ATTENTION`. Diffusers pipelines currently do not support setting different `dtypes` for different stages of the pipeline. In order to run inference in the same way as the original implementation, please refer to the following example. - -The original Mochi implementation zeros out empty prompts. However, enabling this option and placing the entire pipeline under autocast can lead to numerical overflows with the T5 text encoder. +> [!TIP] +> The original Mochi implementation zeros out empty prompts. However, enabling this option and placing the entire pipeline under autocast can lead to numerical overflows with the T5 text encoder. +> +> When enabling `force_zeros_for_empty_prompt`, it is recommended to run the text encoding step outside the autocast context in full precision. -When enabling `force_zeros_for_empty_prompt`, it is recommended to run the text encoding step outside the autocast context in full precision. - - - -Decoding the latents in full precision is very memory intensive. You will need at least 70GB VRAM to generate the 163 frames in this example. To reduce memory, either reduce the number of frames or run the decoding step in `torch.bfloat16`. - +> [!TIP] +> Decoding the latents in full precision is very memory intensive. You will need at least 70GB VRAM to generate the 163 frames in this example. To reduce memory, either reduce the number of frames or run the decoding step in `torch.bfloat16`. ```python import torch @@ -231,9 +229,8 @@ export_to_video(frames, "output.mp4", fps=30) You can use `from_single_file` to load the Mochi transformer in its original format. - -Diffusers currently doesn't support using the FP8 scaled versions of the Mochi single file checkpoints. - +> [!TIP] +> Diffusers currently doesn't support using the FP8 scaled versions of the Mochi single file checkpoints. ```python import torch diff --git a/docs/source/en/api/pipelines/musicldm.md b/docs/source/en/api/pipelines/musicldm.md index c2297162f7..1a83e5932e 100644 --- a/docs/source/en/api/pipelines/musicldm.md +++ b/docs/source/en/api/pipelines/musicldm.md @@ -43,11 +43,8 @@ During inference: * Multiple waveforms can be generated in one go: set `num_waveforms_per_prompt` to a value greater than 1 to enable. Automatic scoring will be performed between the generated waveforms and prompt text, and the audios ranked from best to worst accordingly. * The _length_ of the generated audio sample can be controlled by varying the `audio_length_in_s` argument. - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## MusicLDMPipeline [[autodoc]] MusicLDMPipeline diff --git a/docs/source/en/api/pipelines/omnigen.md b/docs/source/en/api/pipelines/omnigen.md index 074e7b8f01..4fac5c789a 100644 --- a/docs/source/en/api/pipelines/omnigen.md +++ b/docs/source/en/api/pipelines/omnigen.md @@ -21,11 +21,8 @@ The abstract from the paper is: *The emergence of Large Language Models (LLMs) has unified language generation tasks and revolutionized human-machine interaction. However, in the realm of image generation, a unified model capable of handling various tasks within a single framework remains largely unexplored. In this work, we introduce OmniGen, a new diffusion model for unified image generation. OmniGen is characterized by the following features: 1) Unification: OmniGen not only demonstrates text-to-image generation capabilities but also inherently supports various downstream tasks, such as image editing, subject-driven generation, and visual conditional generation. 2) Simplicity: The architecture of OmniGen is highly simplified, eliminating the need for additional plugins. Moreover, compared to existing diffusion models, it is more user-friendly and can complete complex tasks end-to-end through instructions without the need for extra intermediate steps, greatly simplifying the image generation workflow. 3) Knowledge Transfer: Benefit from learning in a unified format, OmniGen effectively transfers knowledge across different tasks, manages unseen tasks and domains, and exhibits novel capabilities. We also explore the model’s reasoning capabilities and potential applications of the chain-of-thought mechanism. This work represents the first attempt at a general-purpose image generation model, and we will release our resources at https://github.com/VectorSpaceLab/OmniGen to foster future advancements.* - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers.md) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading.md#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers.md) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading.md#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. This pipeline was contributed by [staoxiao](https://github.com/staoxiao). The original codebase can be found [here](https://github.com/VectorSpaceLab/OmniGen). The original weights can be found under [hf.co/shitao](https://huggingface.co/Shitao/OmniGen-v1). diff --git a/docs/source/en/api/pipelines/overview.md b/docs/source/en/api/pipelines/overview.md index d3cc318a54..ce883931df 100644 --- a/docs/source/en/api/pipelines/overview.md +++ b/docs/source/en/api/pipelines/overview.md @@ -16,15 +16,12 @@ Pipelines provide a simple way to run state-of-the-art diffusion models in infer All pipelines are built from the base [`DiffusionPipeline`] class which provides basic functionality for loading, downloading, and saving all the components. Specific pipeline types (for example [`StableDiffusionPipeline`]) loaded with [`~DiffusionPipeline.from_pretrained`] are automatically detected and the pipeline components are loaded and passed to the `__init__` function of the pipeline. - - -You shouldn't use the [`DiffusionPipeline`] class for training. Individual components (for example, [`UNet2DModel`] and [`UNet2DConditionModel`]) of diffusion pipelines are usually trained individually, so we suggest directly working with them instead. - -
- -Pipelines do not offer any training functionality. You'll notice PyTorch's autograd is disabled by decorating the [`~DiffusionPipeline.__call__`] method with a [`torch.no_grad`](https://pytorch.org/docs/stable/generated/torch.no_grad.html) decorator because pipelines should not be used for training. If you're interested in training, please take a look at the [Training](../../training/overview) guides instead! - - +> [!WARNING] +> You shouldn't use the [`DiffusionPipeline`] class for training. Individual components (for example, [`UNet2DModel`] and [`UNet2DConditionModel`]) of diffusion pipelines are usually trained individually, so we suggest directly working with them instead. +> +>
+> +> Pipelines do not offer any training functionality. You'll notice PyTorch's autograd is disabled by decorating the [`~DiffusionPipeline.__call__`] method with a [`torch.no_grad`](https://pytorch.org/docs/stable/generated/torch.no_grad.html) decorator because pipelines should not be used for training. If you're interested in training, please take a look at the [Training](../../training/overview) guides instead! The table below lists all the pipelines currently available in 🤗 Diffusers and the tasks they support. Click on a pipeline to view its abstract and published paper. diff --git a/docs/source/en/api/pipelines/pag.md b/docs/source/en/api/pipelines/pag.md index 7b87e58a87..35004b6ad3 100644 --- a/docs/source/en/api/pipelines/pag.md +++ b/docs/source/en/api/pipelines/pag.md @@ -31,11 +31,8 @@ PAG can be used by specifying the `pag_applied_layers` as a parameter when insta - Partial identifier as a RegEx: `down_blocks.2`, or `attn1` - List of identifiers (can be combo of strings and ReGex): `["blocks.1", "blocks.(14|20)", r"down_blocks\.(2,3)"]` - - -Since RegEx is supported as a way for matching layer identifiers, it is crucial to use it correctly otherwise there might be unexpected behaviour. The recommended way to use PAG is by specifying layers as `blocks.{layer_index}` and `blocks.({layer_index_1|layer_index_2|...})`. Using it in any other way, while doable, may bypass our basic validation checks and give you unexpected results. - - +> [!WARNING] +> Since RegEx is supported as a way for matching layer identifiers, it is crucial to use it correctly otherwise there might be unexpected behaviour. The recommended way to use PAG is by specifying layers as `blocks.{layer_index}` and `blocks.({layer_index_1|layer_index_2|...})`. Using it in any other way, while doable, may bypass our basic validation checks and give you unexpected results. ## AnimateDiffPAGPipeline [[autodoc]] AnimateDiffPAGPipeline diff --git a/docs/source/en/api/pipelines/paint_by_example.md b/docs/source/en/api/pipelines/paint_by_example.md index 362c26de68..02bf6db726 100644 --- a/docs/source/en/api/pipelines/paint_by_example.md +++ b/docs/source/en/api/pipelines/paint_by_example.md @@ -27,11 +27,8 @@ The original codebase can be found at [Fantasy-Studio/Paint-by-Example](https:// Paint by Example is supported by the official [Fantasy-Studio/Paint-by-Example](https://huggingface.co/Fantasy-Studio/Paint-by-Example) checkpoint. The checkpoint is warm-started from [CompVis/stable-diffusion-v1-4](https://huggingface.co/CompVis/stable-diffusion-v1-4) to inpaint partly masked images conditioned on example and reference images. - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## PaintByExamplePipeline [[autodoc]] PaintByExamplePipeline diff --git a/docs/source/en/api/pipelines/panorama.md b/docs/source/en/api/pipelines/panorama.md index 9f61388dd5..b65e05dd0b 100644 --- a/docs/source/en/api/pipelines/panorama.md +++ b/docs/source/en/api/pipelines/panorama.md @@ -42,11 +42,8 @@ For example, without circular padding, there is a stitching artifact (default): But with circular padding, the right and the left parts are matching (`circular_padding=True`): ![img](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/indoor_%20circular_padding.png) - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## StableDiffusionPanoramaPipeline [[autodoc]] StableDiffusionPanoramaPipeline diff --git a/docs/source/en/api/pipelines/pia.md b/docs/source/en/api/pipelines/pia.md index 7bd480b49a..eebfa4d4f8 100644 --- a/docs/source/en/api/pipelines/pia.md +++ b/docs/source/en/api/pipelines/pia.md @@ -87,11 +87,8 @@ Here are some sample outputs:
- - -If you plan on using a scheduler that can clip samples, make sure to disable it by setting `clip_sample=False` in the scheduler as this can also have an adverse effect on generated samples. Additionally, the PIA checkpoints can be sensitive to the beta schedule of the scheduler. We recommend setting this to `linear`. - - +> [!TIP] +> If you plan on using a scheduler that can clip samples, make sure to disable it by setting `clip_sample=False` in the scheduler as this can also have an adverse effect on generated samples. Additionally, the PIA checkpoints can be sensitive to the beta schedule of the scheduler. We recommend setting this to `linear`. ## Using FreeInit @@ -149,11 +146,8 @@ export_to_gif(frames, "pia-freeinit-animation.gif") - - -FreeInit is not really free - the improved quality comes at the cost of extra computation. It requires sampling a few extra times depending on the `num_iters` parameter that is set when enabling it. Setting the `use_fast_sampling` parameter to `True` can improve the overall performance (at the cost of lower quality compared to when `use_fast_sampling=False` but still better results than vanilla video generation models). - - +> [!WARNING] +> FreeInit is not really free - the improved quality comes at the cost of extra computation. It requires sampling a few extra times depending on the `num_iters` parameter that is set when enabling it. Setting the `use_fast_sampling` parameter to `True` can improve the overall performance (at the cost of lower quality compared to when `use_fast_sampling=False` but still better results than vanilla video generation models). ## PIAPipeline diff --git a/docs/source/en/api/pipelines/pix2pix.md b/docs/source/en/api/pipelines/pix2pix.md index 20a74577c1..84eb0cb5e5 100644 --- a/docs/source/en/api/pipelines/pix2pix.md +++ b/docs/source/en/api/pipelines/pix2pix.md @@ -24,11 +24,8 @@ The abstract from the paper is: You can find additional information about InstructPix2Pix on the [project page](https://www.timothybrooks.com/instruct-pix2pix), [original codebase](https://github.com/timothybrooks/instruct-pix2pix), and try it out in a [demo](https://huggingface.co/spaces/timbrooks/instruct-pix2pix). - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## StableDiffusionInstructPix2PixPipeline [[autodoc]] StableDiffusionInstructPix2PixPipeline diff --git a/docs/source/en/api/pipelines/pixart.md b/docs/source/en/api/pipelines/pixart.md index a36a2a4b7a..dbdc89857e 100644 --- a/docs/source/en/api/pipelines/pixart.md +++ b/docs/source/en/api/pipelines/pixart.md @@ -29,11 +29,8 @@ Some notes about this pipeline: * It is good at producing high-resolution images at different aspect ratios. To get the best results, the authors recommend some size brackets which can be found [here](https://github.com/PixArt-alpha/PixArt-alpha/blob/08fbbd281ec96866109bdd2cdb75f2f58fb17610/diffusion/data/datasets/utils.py). * It rivals the quality of state-of-the-art text-to-image generation systems (as of this writing) such as Stable Diffusion XL, Imagen, and DALL-E 2, while being more efficient than them. - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## Inference with under 8GB GPU VRAM @@ -112,11 +109,8 @@ del pipe.transformer flush() ``` - - -Notice that while initializing `pipe`, you're setting `text_encoder` to `None` so that it's not loaded. - - +> [!TIP] +> Notice that while initializing `pipe`, you're setting `text_encoder` to `None` so that it's not loaded. Once the latents are computed, pass it off to the VAE to decode into a real image: @@ -133,11 +127,8 @@ By deleting components you aren't using and flushing the GPU VRAM, you should be If you want a report of your memory-usage, run this [script](https://gist.github.com/sayakpaul/3ae0f847001d342af27018a96f467e4e). - - -Text embeddings computed in 8-bit can impact the quality of the generated images because of the information loss in the representation space caused by the reduced precision. It's recommended to compare the outputs with and without 8-bit. - - +> [!WARNING] +> Text embeddings computed in 8-bit can impact the quality of the generated images because of the information loss in the representation space caused by the reduced precision. It's recommended to compare the outputs with and without 8-bit. While loading the `text_encoder`, you set `load_in_8bit` to `True`. You could also specify `load_in_4bit` to bring your memory requirements down even further to under 7GB. diff --git a/docs/source/en/api/pipelines/pixart_sigma.md b/docs/source/en/api/pipelines/pixart_sigma.md index dded4ea2d7..06b54de43b 100644 --- a/docs/source/en/api/pipelines/pixart_sigma.md +++ b/docs/source/en/api/pipelines/pixart_sigma.md @@ -31,17 +31,11 @@ Some notes about this pipeline: * It shows the ability of generating super high resolution images, such as 2048px or even 4K. * It shows that text-to-image models can grow from a weak model to a stronger one through several improvements (VAEs, datasets, and so on.) - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - - - - -You can further improve generation quality by passing the generated image from [`PixArtSigmaPipeline`] to the [SDXL refiner](../../using-diffusers/sdxl#base-to-refiner-model) model. - - +> [!TIP] +> You can further improve generation quality by passing the generated image from [`PixArtSigmaPipeline`] to the [SDXL refiner](../../using-diffusers/sdxl#base-to-refiner-model) model. ## Inference with under 8GB GPU VRAM @@ -119,11 +113,8 @@ del pipe.transformer flush() ``` - - -Notice that while initializing `pipe`, you're setting `text_encoder` to `None` so that it's not loaded. - - +> [!TIP] +> Notice that while initializing `pipe`, you're setting `text_encoder` to `None` so that it's not loaded. Once the latents are computed, pass it off to the VAE to decode into a real image: @@ -140,11 +131,8 @@ By deleting components you aren't using and flushing the GPU VRAM, you should be If you want a report of your memory-usage, run this [script](https://gist.github.com/sayakpaul/3ae0f847001d342af27018a96f467e4e). - - -Text embeddings computed in 8-bit can impact the quality of the generated images because of the information loss in the representation space caused by the reduced precision. It's recommended to compare the outputs with and without 8-bit. - - +> [!WARNING] +> Text embeddings computed in 8-bit can impact the quality of the generated images because of the information loss in the representation space caused by the reduced precision. It's recommended to compare the outputs with and without 8-bit. While loading the `text_encoder`, you set `load_in_8bit` to `True`. You could also specify `load_in_4bit` to bring your memory requirements down even further to under 7GB. diff --git a/docs/source/en/api/pipelines/qwenimage.md b/docs/source/en/api/pipelines/qwenimage.md index 4c999bca35..27cc4802f6 100644 --- a/docs/source/en/api/pipelines/qwenimage.md +++ b/docs/source/en/api/pipelines/qwenimage.md @@ -28,11 +28,8 @@ Qwen-Image comes in the following variants: | Qwen-Image-Edit | [`Qwen/Qwen-Image-Edit`](https://huggingface.co/Qwen/Qwen-Image-Edit) | | Qwen-Image-Edit Plus | [Qwen/Qwen-Image-Edit-2509](https://huggingface.co/Qwen/Qwen-Image-Edit-2509) | - - -[Caching](../../optimization/cache) may also speed up inference by storing and reusing intermediate outputs. - - +> [!TIP] +> [Caching](../../optimization/cache) may also speed up inference by storing and reusing intermediate outputs. ## LoRA for faster inference @@ -91,11 +88,8 @@ image.save("qwen_fewsteps.png") - - -The `guidance_scale` parameter in the pipeline is there to support future guidance-distilled models when they come up. Note that passing `guidance_scale` to the pipeline is ineffective. To enable classifier-free guidance, please pass `true_cfg_scale` and `negative_prompt` (even an empty negative prompt like " ") should enable classifier-free guidance computations. - - +> [!TIP] +> The `guidance_scale` parameter in the pipeline is there to support future guidance-distilled models when they come up. Note that passing `guidance_scale` to the pipeline is ineffective. To enable classifier-free guidance, please pass `true_cfg_scale` and `negative_prompt` (even an empty negative prompt like " ") should enable classifier-free guidance computations. ## Multi-image reference with QwenImageEditPlusPipeline diff --git a/docs/source/en/api/pipelines/sana.md b/docs/source/en/api/pipelines/sana.md index 7491689fd8..a948620f96 100644 --- a/docs/source/en/api/pipelines/sana.md +++ b/docs/source/en/api/pipelines/sana.md @@ -25,11 +25,8 @@ The abstract from the paper is: *We introduce Sana, a text-to-image framework that can efficiently generate images up to 4096×4096 resolution. Sana can synthesize high-resolution, high-quality images with strong text-image alignment at a remarkably fast speed, deployable on laptop GPU. Core designs include: (1) Deep compression autoencoder: unlike traditional AEs, which compress images only 8×, we trained an AE that can compress images 32×, effectively reducing the number of latent tokens. (2) Linear DiT: we replace all vanilla attention in DiT with linear attention, which is more efficient at high resolutions without sacrificing quality. (3) Decoder-only text encoder: we replaced T5 with modern decoder-only small LLM as the text encoder and designed complex human instruction with in-context learning to enhance the image-text alignment. (4) Efficient training and sampling: we propose Flow-DPM-Solver to reduce sampling steps, with efficient caption labeling and selection to accelerate convergence. As a result, Sana-0.6B is very competitive with modern giant diffusion model (e.g. Flux-12B), being 20 times smaller and 100+ times faster in measured throughput. Moreover, Sana-0.6B can be deployed on a 16GB laptop GPU, taking less than 1 second to generate a 1024×1024 resolution image. Sana enables content creation at low cost. Code and model will be publicly released.* - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. This pipeline was contributed by [lawrence-cj](https://github.com/lawrence-cj) and [chenjy2003](https://github.com/chenjy2003). The original codebase can be found [here](https://github.com/NVlabs/Sana). The original weights can be found under [hf.co/Efficient-Large-Model](https://huggingface.co/Efficient-Large-Model). @@ -49,11 +46,8 @@ Refer to [this](https://huggingface.co/collections/Efficient-Large-Model/sana-67 Note: The recommended dtype mentioned is for the transformer weights. The text encoder and VAE weights must stay in `torch.bfloat16` or `torch.float32` for the model to work correctly. Please refer to the inference example below to see how to load the model with the recommended dtype. - - -Make sure to pass the `variant` argument for downloaded checkpoints to use lower disk space. Set it to `"fp16"` for models with recommended dtype as `torch.float16`, and `"bf16"` for models with recommended dtype as `torch.bfloat16`. By default, `torch.float32` weights are downloaded, which use twice the amount of disk storage. Additionally, `torch.float32` weights can be downcasted on-the-fly by specifying the `torch_dtype` argument. Read about it in the [docs](https://huggingface.co/docs/diffusers/v0.31.0/en/api/pipelines/overview#diffusers.DiffusionPipeline.from_pretrained). - - +> [!TIP] +> Make sure to pass the `variant` argument for downloaded checkpoints to use lower disk space. Set it to `"fp16"` for models with recommended dtype as `torch.float16`, and `"bf16"` for models with recommended dtype as `torch.bfloat16`. By default, `torch.float32` weights are downloaded, which use twice the amount of disk storage. Additionally, `torch.float32` weights can be downcasted on-the-fly by specifying the `torch_dtype` argument. Read about it in the [docs](https://huggingface.co/docs/diffusers/v0.31.0/en/api/pipelines/overview#diffusers.DiffusionPipeline.from_pretrained). ## Quantization diff --git a/docs/source/en/api/pipelines/sana_sprint.md b/docs/source/en/api/pipelines/sana_sprint.md index 93ab9fe418..357d7e406d 100644 --- a/docs/source/en/api/pipelines/sana_sprint.md +++ b/docs/source/en/api/pipelines/sana_sprint.md @@ -24,11 +24,8 @@ The abstract from the paper is: *This paper presents SANA-Sprint, an efficient diffusion model for ultra-fast text-to-image (T2I) generation. SANA-Sprint is built on a pre-trained foundation model and augmented with hybrid distillation, dramatically reducing inference steps from 20 to 1-4. We introduce three key innovations: (1) We propose a training-free approach that transforms a pre-trained flow-matching model for continuous-time consistency distillation (sCM), eliminating costly training from scratch and achieving high training efficiency. Our hybrid distillation strategy combines sCM with latent adversarial distillation (LADD): sCM ensures alignment with the teacher model, while LADD enhances single-step generation fidelity. (2) SANA-Sprint is a unified step-adaptive model that achieves high-quality generation in 1-4 steps, eliminating step-specific training and improving efficiency. (3) We integrate ControlNet with SANA-Sprint for real-time interactive image generation, enabling instant visual feedback for user interaction. SANA-Sprint establishes a new Pareto frontier in speed-quality tradeoffs, achieving state-of-the-art performance with 7.59 FID and 0.74 GenEval in only 1 step — outperforming FLUX-schnell (7.94 FID / 0.71 GenEval) while being 10× faster (0.1s vs 1.1s on H100). It also achieves 0.1s (T2I) and 0.25s (ControlNet) latency for 1024×1024 images on H100, and 0.31s (T2I) on an RTX 4090, showcasing its exceptional efficiency and potential for AI-powered consumer applications (AIPC). Code and pre-trained models will be open-sourced.* - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. This pipeline was contributed by [lawrence-cj](https://github.com/lawrence-cj), [shuchen Xue](https://github.com/scxue) and [Enze Xie](https://github.com/xieenze). The original codebase can be found [here](https://github.com/NVlabs/Sana). The original weights can be found under [hf.co/Efficient-Large-Model](https://huggingface.co/Efficient-Large-Model/). diff --git a/docs/source/en/api/pipelines/self_attention_guidance.md b/docs/source/en/api/pipelines/self_attention_guidance.md index 5578fdfa63..8d411598ae 100644 --- a/docs/source/en/api/pipelines/self_attention_guidance.md +++ b/docs/source/en/api/pipelines/self_attention_guidance.md @@ -23,11 +23,8 @@ The abstract from the paper is: You can find additional information about Self-Attention Guidance on the [project page](https://ku-cvlab.github.io/Self-Attention-Guidance), [original codebase](https://github.com/KU-CVLAB/Self-Attention-Guidance), and try it out in a [demo](https://huggingface.co/spaces/susunghong/Self-Attention-Guidance) or [notebook](https://colab.research.google.com/github/SusungHong/Self-Attention-Guidance/blob/main/SAG_Stable.ipynb). - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## StableDiffusionSAGPipeline [[autodoc]] StableDiffusionSAGPipeline diff --git a/docs/source/en/api/pipelines/semantic_stable_diffusion.md b/docs/source/en/api/pipelines/semantic_stable_diffusion.md index 1ce44cf2de..dda428e80f 100644 --- a/docs/source/en/api/pipelines/semantic_stable_diffusion.md +++ b/docs/source/en/api/pipelines/semantic_stable_diffusion.md @@ -22,11 +22,8 @@ The abstract from the paper is: *Text-to-image diffusion models have recently received a lot of interest for their astonishing ability to produce high-fidelity images from text only. However, achieving one-shot generation that aligns with the user's intent is nearly impossible, yet small changes to the input prompt often result in very different images. This leaves the user with little semantic control. To put the user in control, we show how to interact with the diffusion process to flexibly steer it along semantic directions. This semantic guidance (SEGA) generalizes to any generative architecture using classifier-free guidance. More importantly, it allows for subtle and extensive edits, changes in composition and style, as well as optimizing the overall artistic conception. We demonstrate SEGA's effectiveness on both latent and pixel-based diffusion models such as Stable Diffusion, Paella, and DeepFloyd-IF using a variety of tasks, thus providing strong evidence for its versatility, flexibility, and improvements over existing methods.* - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## SemanticStableDiffusionPipeline [[autodoc]] SemanticStableDiffusionPipeline diff --git a/docs/source/en/api/pipelines/shap_e.md b/docs/source/en/api/pipelines/shap_e.md index 5e5af0656a..3e505894ca 100644 --- a/docs/source/en/api/pipelines/shap_e.md +++ b/docs/source/en/api/pipelines/shap_e.md @@ -17,11 +17,8 @@ The abstract from the paper is: The original codebase can be found at [openai/shap-e](https://github.com/openai/shap-e). - - -See the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> See the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## ShapEPipeline [[autodoc]] ShapEPipeline diff --git a/docs/source/en/api/pipelines/stable_cascade.md b/docs/source/en/api/pipelines/stable_cascade.md index b47f9ae337..70de6776e9 100644 --- a/docs/source/en/api/pipelines/stable_cascade.md +++ b/docs/source/en/api/pipelines/stable_cascade.md @@ -41,15 +41,12 @@ The Stage C model operates on the small 24 x 24 latents and denoises the latents The Stage B and Stage A models are used with the `StableCascadeDecoderPipeline` and are responsible for generating the final image given the small 24 x 24 latents. - - -There are some restrictions on data types that can be used with the Stable Cascade models. The official checkpoints for the `StableCascadePriorPipeline` do not support the `torch.float16` data type. Please use `torch.bfloat16` instead. - -In order to use the `torch.bfloat16` data type with the `StableCascadeDecoderPipeline` you need to have PyTorch 2.2.0 or higher installed. This also means that using the `StableCascadeCombinedPipeline` with `torch.bfloat16` requires PyTorch 2.2.0 or higher, since it calls the `StableCascadeDecoderPipeline` internally. - -If it is not possible to install PyTorch 2.2.0 or higher in your environment, the `StableCascadeDecoderPipeline` can be used on its own with the `torch.float16` data type. You can download the full precision or `bf16` variant weights for the pipeline and cast the weights to `torch.float16`. - - +> [!WARNING] +> There are some restrictions on data types that can be used with the Stable Cascade models. The official checkpoints for the `StableCascadePriorPipeline` do not support the `torch.float16` data type. Please use `torch.bfloat16` instead. +> +> In order to use the `torch.bfloat16` data type with the `StableCascadeDecoderPipeline` you need to have PyTorch 2.2.0 or higher installed. This also means that using the `StableCascadeCombinedPipeline` with `torch.bfloat16` requires PyTorch 2.2.0 or higher, since it calls the `StableCascadeDecoderPipeline` internally. +> +> If it is not possible to install PyTorch 2.2.0 or higher in your environment, the `StableCascadeDecoderPipeline` can be used on its own with the `torch.float16` data type. You can download the full precision or `bf16` variant weights for the pipeline and cast the weights to `torch.float16`. ## Usage example diff --git a/docs/source/en/api/pipelines/stable_diffusion/depth2img.md b/docs/source/en/api/pipelines/stable_diffusion/depth2img.md index e198eaa952..aa43cf7db9 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/depth2img.md +++ b/docs/source/en/api/pipelines/stable_diffusion/depth2img.md @@ -18,13 +18,10 @@ specific language governing permissions and limitations under the License. The Stable Diffusion model can also infer depth based on an image using [MiDaS](https://github.com/isl-org/MiDaS). This allows you to pass a text prompt and an initial image to condition the generation of new images as well as a `depth_map` to preserve the image structure. - - -Make sure to check out the Stable Diffusion [Tips](overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality, and how to reuse pipeline components efficiently! - -If you're interested in using one of the official checkpoints for a task, explore the [CompVis](https://huggingface.co/CompVis), [Runway](https://huggingface.co/runwayml), and [Stability AI](https://huggingface.co/stabilityai) Hub organizations! - - +> [!TIP] +> Make sure to check out the Stable Diffusion [Tips](overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality, and how to reuse pipeline components efficiently! +> +> If you're interested in using one of the official checkpoints for a task, explore the [CompVis](https://huggingface.co/CompVis), [Runway](https://huggingface.co/runwayml), and [Stability AI](https://huggingface.co/stabilityai) Hub organizations! ## StableDiffusionDepth2ImgPipeline diff --git a/docs/source/en/api/pipelines/stable_diffusion/gligen.md b/docs/source/en/api/pipelines/stable_diffusion/gligen.md index e9704fc1de..c8297fb7b3 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/gligen.md +++ b/docs/source/en/api/pipelines/stable_diffusion/gligen.md @@ -21,13 +21,10 @@ The abstract from the [paper](https://huggingface.co/papers/2301.07093) is: *Large-scale text-to-image diffusion models have made amazing advances. However, the status quo is to use text input alone, which can impede controllability. In this work, we propose GLIGEN, Grounded-Language-to-Image Generation, a novel approach that builds upon and extends the functionality of existing pre-trained text-to-image diffusion models by enabling them to also be conditioned on grounding inputs. To preserve the vast concept knowledge of the pre-trained model, we freeze all of its weights and inject the grounding information into new trainable layers via a gated mechanism. Our model achieves open-world grounded text2img generation with caption and bounding box condition inputs, and the grounding ability generalizes well to novel spatial configurations and concepts. GLIGEN’s zeroshot performance on COCO and LVIS outperforms existing supervised layout-to-image baselines by a large margin.* - - -Make sure to check out the Stable Diffusion [Tips](https://huggingface.co/docs/diffusers/en/api/pipelines/stable_diffusion/overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality and how to reuse pipeline components efficiently! - -If you want to use one of the official checkpoints for a task, explore the [gligen](https://huggingface.co/gligen) Hub organizations! - - +> [!TIP] +> Make sure to check out the Stable Diffusion [Tips](https://huggingface.co/docs/diffusers/en/api/pipelines/stable_diffusion/overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality and how to reuse pipeline components efficiently! +> +> If you want to use one of the official checkpoints for a task, explore the [gligen](https://huggingface.co/gligen) Hub organizations! [`StableDiffusionGLIGENPipeline`] was contributed by [Nikhil Gajendrakumar](https://github.com/nikhil-masterful) and [`StableDiffusionGLIGENTextImagePipeline`] was contributed by [Nguyễn Công Tú Anh](https://github.com/tuanh123789). diff --git a/docs/source/en/api/pipelines/stable_diffusion/image_variation.md b/docs/source/en/api/pipelines/stable_diffusion/image_variation.md index 7a50971fdf..b1b7146b33 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/image_variation.md +++ b/docs/source/en/api/pipelines/stable_diffusion/image_variation.md @@ -16,11 +16,8 @@ The Stable Diffusion model can also generate variations from an input image. It The original codebase can be found at [LambdaLabsML/lambda-diffusers](https://github.com/LambdaLabsML/lambda-diffusers#stable-diffusion-image-variations) and additional official checkpoints for image variation can be found at [lambdalabs/sd-image-variations-diffusers](https://huggingface.co/lambdalabs/sd-image-variations-diffusers). - - -Make sure to check out the Stable Diffusion [Tips](./overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality, and how to reuse pipeline components efficiently! - - +> [!TIP] +> Make sure to check out the Stable Diffusion [Tips](./overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality, and how to reuse pipeline components efficiently! ## StableDiffusionImageVariationPipeline diff --git a/docs/source/en/api/pipelines/stable_diffusion/img2img.md b/docs/source/en/api/pipelines/stable_diffusion/img2img.md index bec67b4f4e..f9e4476427 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/img2img.md +++ b/docs/source/en/api/pipelines/stable_diffusion/img2img.md @@ -24,11 +24,8 @@ The abstract from the paper is: *Guided image synthesis enables everyday users to create and edit photo-realistic images with minimum effort. The key challenge is balancing faithfulness to the user input (e.g., hand-drawn colored strokes) and realism of the synthesized image. Existing GAN-based methods attempt to achieve such balance using either conditional GANs or GAN inversions, which are challenging and often require additional training data or loss functions for individual applications. To address these issues, we introduce a new image synthesis and editing method, Stochastic Differential Editing (SDEdit), based on a diffusion model generative prior, which synthesizes realistic images by iteratively denoising through a stochastic differential equation (SDE). Given an input image with user guide of any type, SDEdit first adds noise to the input, then subsequently denoises the resulting image through the SDE prior to increase its realism. SDEdit does not require task-specific training or inversions and can naturally achieve the balance between realism and faithfulness. SDEdit significantly outperforms state-of-the-art GAN-based methods by up to 98.09% on realism and 91.72% on overall satisfaction scores, according to a human perception study, on multiple tasks, including stroke-based image synthesis and editing as well as image compositing.* - - -Make sure to check out the Stable Diffusion [Tips](overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality, and how to reuse pipeline components efficiently! - - +> [!TIP] +> Make sure to check out the Stable Diffusion [Tips](overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality, and how to reuse pipeline components efficiently! ## StableDiffusionImg2ImgPipeline diff --git a/docs/source/en/api/pipelines/stable_diffusion/inpaint.md b/docs/source/en/api/pipelines/stable_diffusion/inpaint.md index 0b558b2fc0..84cc31e158 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/inpaint.md +++ b/docs/source/en/api/pipelines/stable_diffusion/inpaint.md @@ -25,13 +25,10 @@ as [runwayml/stable-diffusion-inpainting](https://huggingface.co/runwayml/stable text-to-image Stable Diffusion checkpoints, such as [stable-diffusion-v1-5/stable-diffusion-v1-5](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) are also compatible but they might be less performant. - - -Make sure to check out the Stable Diffusion [Tips](overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality, and how to reuse pipeline components efficiently! - -If you're interested in using one of the official checkpoints for a task, explore the [CompVis](https://huggingface.co/CompVis), [Runway](https://huggingface.co/runwayml), and [Stability AI](https://huggingface.co/stabilityai) Hub organizations! - - +> [!TIP] +> Make sure to check out the Stable Diffusion [Tips](overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality, and how to reuse pipeline components efficiently! +> +> If you're interested in using one of the official checkpoints for a task, explore the [CompVis](https://huggingface.co/CompVis), [Runway](https://huggingface.co/runwayml), and [Stability AI](https://huggingface.co/stabilityai) Hub organizations! ## StableDiffusionInpaintPipeline diff --git a/docs/source/en/api/pipelines/stable_diffusion/latent_upscale.md b/docs/source/en/api/pipelines/stable_diffusion/latent_upscale.md index d5a15cb002..4f0521740c 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/latent_upscale.md +++ b/docs/source/en/api/pipelines/stable_diffusion/latent_upscale.md @@ -14,13 +14,10 @@ specific language governing permissions and limitations under the License. The Stable Diffusion latent upscaler model was created by [Katherine Crowson](https://github.com/crowsonkb/k-diffusion) in collaboration with [Stability AI](https://stability.ai/). It is used to enhance the output image resolution by a factor of 2 (see this demo [notebook](https://colab.research.google.com/drive/1o1qYJcFeywzCIdkfKJy7cTpgZTCM2EI4) for a demonstration of the original implementation). - - -Make sure to check out the Stable Diffusion [Tips](overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality, and how to reuse pipeline components efficiently! - -If you're interested in using one of the official checkpoints for a task, explore the [CompVis](https://huggingface.co/CompVis), [Runway](https://huggingface.co/runwayml), and [Stability AI](https://huggingface.co/stabilityai) Hub organizations! - - +> [!TIP] +> Make sure to check out the Stable Diffusion [Tips](overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality, and how to reuse pipeline components efficiently! +> +> If you're interested in using one of the official checkpoints for a task, explore the [CompVis](https://huggingface.co/CompVis), [Runway](https://huggingface.co/runwayml), and [Stability AI](https://huggingface.co/stabilityai) Hub organizations! ## StableDiffusionLatentUpscalePipeline diff --git a/docs/source/en/api/pipelines/stable_diffusion/ldm3d_diffusion.md b/docs/source/en/api/pipelines/stable_diffusion/ldm3d_diffusion.md index 4c52ed90f0..15f9f1db85 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/ldm3d_diffusion.md +++ b/docs/source/en/api/pipelines/stable_diffusion/ldm3d_diffusion.md @@ -30,11 +30,8 @@ The abstract from the paper is: *This research paper proposes a Latent Diffusion Model for 3D (LDM3D) that generates both image and depth map data from a given text prompt, allowing users to generate RGBD images from text prompts. The LDM3D model is fine-tuned on a dataset of tuples containing an RGB image, depth map and caption, and validated through extensive experiments. We also develop an application called DepthFusion, which uses the generated RGB images and depth maps to create immersive and interactive 360-degree-view experiences using TouchDesigner. This technology has the potential to transform a wide range of industries, from entertainment and gaming to architecture and design. Overall, this paper presents a significant contribution to the field of generative AI and computer vision, and showcases the potential of LDM3D and DepthFusion to revolutionize content creation and digital experiences. A short video summarizing the approach can be found at [this url](https://t.ly/tdi2).* - - -Make sure to check out the Stable Diffusion [Tips](overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality, and how to reuse pipeline components efficiently! - - +> [!TIP] +> Make sure to check out the Stable Diffusion [Tips](overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality, and how to reuse pipeline components efficiently! ## StableDiffusionLDM3DPipeline diff --git a/docs/source/en/api/pipelines/stable_diffusion/sdxl_turbo.md b/docs/source/en/api/pipelines/stable_diffusion/sdxl_turbo.md index aac3a3d870..7964db4c9d 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/sdxl_turbo.md +++ b/docs/source/en/api/pipelines/stable_diffusion/sdxl_turbo.md @@ -26,10 +26,7 @@ The abstract from the paper is: - SDXL Turbo has been trained to generate images of size 512x512. - SDXL Turbo is open-access, but not open-source meaning that one might have to buy a model license in order to use it for commercial applications. Make sure to read the [official model card](https://huggingface.co/stabilityai/sdxl-turbo) to learn more. - - -To learn how to use SDXL Turbo for various tasks, how to optimize performance, and other usage examples, take a look at the [SDXL Turbo](../../../using-diffusers/sdxl_turbo) guide. - -Check out the [Stability AI](https://huggingface.co/stabilityai) Hub organization for the official base and refiner model checkpoints! - - +> [!TIP] +> To learn how to use SDXL Turbo for various tasks, how to optimize performance, and other usage examples, take a look at the [SDXL Turbo](../../../using-diffusers/sdxl_turbo) guide. +> +> Check out the [Stability AI](https://huggingface.co/stabilityai) Hub organization for the official base and refiner model checkpoints! diff --git a/docs/source/en/api/pipelines/stable_diffusion/stable_diffusion_2.md b/docs/source/en/api/pipelines/stable_diffusion/stable_diffusion_2.md index 89e9f5305e..67729cd195 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/stable_diffusion_2.md +++ b/docs/source/en/api/pipelines/stable_diffusion/stable_diffusion_2.md @@ -33,13 +33,10 @@ Stable Diffusion 2 is available for tasks like text-to-image, inpainting, super- Here are some examples for how to use Stable Diffusion 2 for each task: - - -Make sure to check out the Stable Diffusion [Tips](overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality, and how to reuse pipeline components efficiently! - -If you're interested in using one of the official checkpoints for a task, explore the [CompVis](https://huggingface.co/CompVis), [Runway](https://huggingface.co/runwayml), and [Stability AI](https://huggingface.co/stabilityai) Hub organizations! - - +> [!TIP] +> Make sure to check out the Stable Diffusion [Tips](overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality, and how to reuse pipeline components efficiently! +> +> If you're interested in using one of the official checkpoints for a task, explore the [CompVis](https://huggingface.co/CompVis), [Runway](https://huggingface.co/runwayml), and [Stability AI](https://huggingface.co/stabilityai) Hub organizations! ## Text-to-image diff --git a/docs/source/en/api/pipelines/stable_diffusion/stable_diffusion_3.md b/docs/source/en/api/pipelines/stable_diffusion/stable_diffusion_3.md index 211b26889a..3c49df101c 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/stable_diffusion_3.md +++ b/docs/source/en/api/pipelines/stable_diffusion/stable_diffusion_3.md @@ -34,11 +34,8 @@ Use the command below to log in: hf auth login ``` - - -The SD3 pipeline uses three text encoders to generate an image. Model offloading is necessary in order for it to run on most commodity hardware. Please use the `torch.float16` data type for additional memory savings. - - +> [!TIP] +> The SD3 pipeline uses three text encoders to generate an image. Model offloading is necessary in order for it to run on most commodity hardware. Please use the `torch.float16` data type for additional memory savings. ```python import torch @@ -124,11 +121,8 @@ image.save("result.jpg") - - -Check out [IP-Adapter](../../../using-diffusers/ip_adapter) to learn more about how IP-Adapters work. - - +> [!TIP] +> Check out [IP-Adapter](../../../using-diffusers/ip_adapter) to learn more about how IP-Adapters work. ## Memory Optimisations for SD3 @@ -333,11 +327,8 @@ image = pipe( You can send a different prompt to the CLIP Text Encoders and the T5 Text Encoder to prevent the prompt from being truncated by the CLIP Text Encoders and to improve generation. - - -The prompt with the CLIP Text Encoders is still truncated to the 77 token limit. - - +> [!TIP] +> The prompt with the CLIP Text Encoders is still truncated to the 77 token limit. ```python prompt = "A whimsical and creative image depicting a hybrid creature that is a mix of a waffle and a hippopotamus, basking in a river of melted butter amidst a breakfast-themed landscape. A river of warm, melted butter, pancake-like foliage in the background, a towering pepper mill standing in for a tree." diff --git a/docs/source/en/api/pipelines/stable_diffusion/stable_diffusion_safe.md b/docs/source/en/api/pipelines/stable_diffusion/stable_diffusion_safe.md index 1736491107..151b0b8a65 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/stable_diffusion_safe.md +++ b/docs/source/en/api/pipelines/stable_diffusion/stable_diffusion_safe.md @@ -45,11 +45,8 @@ There are 4 configurations (`SafetyConfig.WEAK`, `SafetyConfig.MEDIUM`, `SafetyC >>> out = pipeline(prompt=prompt, **SafetyConfig.MAX) ``` - - -Make sure to check out the Stable Diffusion [Tips](overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality, and how to reuse pipeline components efficiently! - - +> [!TIP] +> Make sure to check out the Stable Diffusion [Tips](overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality, and how to reuse pipeline components efficiently! ## StableDiffusionPipelineSafe diff --git a/docs/source/en/api/pipelines/stable_diffusion/stable_diffusion_xl.md b/docs/source/en/api/pipelines/stable_diffusion/stable_diffusion_xl.md index 30e4379066..6863d408b5 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/stable_diffusion_xl.md +++ b/docs/source/en/api/pipelines/stable_diffusion/stable_diffusion_xl.md @@ -33,13 +33,10 @@ The abstract from the paper is: - SDXL output images can be improved by making use of a refiner model in an image-to-image setting. - SDXL offers `negative_original_size`, `negative_crops_coords_top_left`, and `negative_target_size` to negatively condition the model on image resolution and cropping parameters. - - -To learn how to use SDXL for various tasks, how to optimize performance, and other usage examples, take a look at the [Stable Diffusion XL](../../../using-diffusers/sdxl) guide. - -Check out the [Stability AI](https://huggingface.co/stabilityai) Hub organization for the official base and refiner model checkpoints! - - +> [!TIP] +> To learn how to use SDXL for various tasks, how to optimize performance, and other usage examples, take a look at the [Stable Diffusion XL](../../../using-diffusers/sdxl) guide. +> +> Check out the [Stability AI](https://huggingface.co/stabilityai) Hub organization for the official base and refiner model checkpoints! ## StableDiffusionXLPipeline diff --git a/docs/source/en/api/pipelines/stable_diffusion/svd.md b/docs/source/en/api/pipelines/stable_diffusion/svd.md index ab51f9b663..0c33c06007 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/svd.md +++ b/docs/source/en/api/pipelines/stable_diffusion/svd.md @@ -18,15 +18,12 @@ The abstract from the paper is: *We present Stable Video Diffusion - a latent video diffusion model for high-resolution, state-of-the-art text-to-video and image-to-video generation. Recently, latent diffusion models trained for 2D image synthesis have been turned into generative video models by inserting temporal layers and finetuning them on small, high-quality video datasets. However, training methods in the literature vary widely, and the field has yet to agree on a unified strategy for curating video data. In this paper, we identify and evaluate three different stages for successful training of video LDMs: text-to-image pretraining, video pretraining, and high-quality video finetuning. Furthermore, we demonstrate the necessity of a well-curated pretraining dataset for generating high-quality videos and present a systematic curation process to train a strong base model, including captioning and filtering strategies. We then explore the impact of finetuning our base model on high-quality data and train a text-to-video model that is competitive with closed-source video generation. We also show that our base model provides a powerful motion representation for downstream tasks such as image-to-video generation and adaptability to camera motion-specific LoRA modules. Finally, we demonstrate that our model provides a strong multi-view 3D-prior and can serve as a base to finetune a multi-view diffusion model that jointly generates multiple views of objects in a feedforward fashion, outperforming image-based methods at a fraction of their compute budget. We release code and model weights at this https URL.* - - -To learn how to use Stable Video Diffusion, take a look at the [Stable Video Diffusion](../../../using-diffusers/svd) guide. - -
- -Check out the [Stability AI](https://huggingface.co/stabilityai) Hub organization for the [base](https://huggingface.co/stabilityai/stable-video-diffusion-img2vid) and [extended frame](https://huggingface.co/stabilityai/stable-video-diffusion-img2vid-xt) checkpoints! - - +> [!TIP] +> To learn how to use Stable Video Diffusion, take a look at the [Stable Video Diffusion](../../../using-diffusers/svd) guide. +> +>
+> +> Check out the [Stability AI](https://huggingface.co/stabilityai) Hub organization for the [base](https://huggingface.co/stabilityai/stable-video-diffusion-img2vid) and [extended frame](https://huggingface.co/stabilityai/stable-video-diffusion-img2vid-xt) checkpoints! ## Tips diff --git a/docs/source/en/api/pipelines/stable_diffusion/text2img.md b/docs/source/en/api/pipelines/stable_diffusion/text2img.md index c17348c8ff..59a0f00d22 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/text2img.md +++ b/docs/source/en/api/pipelines/stable_diffusion/text2img.md @@ -22,13 +22,10 @@ The abstract from the paper is: *By decomposing the image formation process into a sequential application of denoising autoencoders, diffusion models (DMs) achieve state-of-the-art synthesis results on image data and beyond. Additionally, their formulation allows for a guiding mechanism to control the image generation process without retraining. However, since these models typically operate directly in pixel space, optimization of powerful DMs often consumes hundreds of GPU days and inference is expensive due to sequential evaluations. To enable DM training on limited computational resources while retaining their quality and flexibility, we apply them in the latent space of powerful pretrained autoencoders. In contrast to previous work, training diffusion models on such a representation allows for the first time to reach a near-optimal point between complexity reduction and detail preservation, greatly boosting visual fidelity. By introducing cross-attention layers into the model architecture, we turn diffusion models into powerful and flexible generators for general conditioning inputs such as text or bounding boxes and high-resolution synthesis becomes possible in a convolutional manner. Our latent diffusion models (LDMs) achieve a new state of the art for image inpainting and highly competitive performance on various tasks, including unconditional image generation, semantic scene synthesis, and super-resolution, while significantly reducing computational requirements compared to pixel-based DMs. Code is available at https://github.com/CompVis/latent-diffusion.* - - -Make sure to check out the Stable Diffusion [Tips](overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality, and how to reuse pipeline components efficiently! - -If you're interested in using one of the official checkpoints for a task, explore the [CompVis](https://huggingface.co/CompVis), [Runway](https://huggingface.co/runwayml), and [Stability AI](https://huggingface.co/stabilityai) Hub organizations! - - +> [!TIP] +> Make sure to check out the Stable Diffusion [Tips](overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality, and how to reuse pipeline components efficiently! +> +> If you're interested in using one of the official checkpoints for a task, explore the [CompVis](https://huggingface.co/CompVis), [Runway](https://huggingface.co/runwayml), and [Stability AI](https://huggingface.co/stabilityai) Hub organizations! ## StableDiffusionPipeline diff --git a/docs/source/en/api/pipelines/stable_diffusion/upscale.md b/docs/source/en/api/pipelines/stable_diffusion/upscale.md index 411491263c..14393370be 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/upscale.md +++ b/docs/source/en/api/pipelines/stable_diffusion/upscale.md @@ -18,13 +18,10 @@ specific language governing permissions and limitations under the License. The Stable Diffusion upscaler diffusion model was created by the researchers and engineers from [CompVis](https://github.com/CompVis), [Stability AI](https://stability.ai/), and [LAION](https://laion.ai/). It is used to enhance the resolution of input images by a factor of 4. - - -Make sure to check out the Stable Diffusion [Tips](overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality, and how to reuse pipeline components efficiently! - -If you're interested in using one of the official checkpoints for a task, explore the [CompVis](https://huggingface.co/CompVis), [Runway](https://huggingface.co/runwayml), and [Stability AI](https://huggingface.co/stabilityai) Hub organizations! - - +> [!TIP] +> Make sure to check out the Stable Diffusion [Tips](overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality, and how to reuse pipeline components efficiently! +> +> If you're interested in using one of the official checkpoints for a task, explore the [CompVis](https://huggingface.co/CompVis), [Runway](https://huggingface.co/runwayml), and [Stability AI](https://huggingface.co/stabilityai) Hub organizations! ## StableDiffusionUpscalePipeline diff --git a/docs/source/en/api/pipelines/stable_unclip.md b/docs/source/en/api/pipelines/stable_unclip.md index 5abb6028c4..09100201bb 100644 --- a/docs/source/en/api/pipelines/stable_unclip.md +++ b/docs/source/en/api/pipelines/stable_unclip.md @@ -65,11 +65,8 @@ wave_prompt = "dramatic wave, the Oceans roar, Strong wave spiral across the oce image = pipe(prompt=wave_prompt).images[0] image ``` - - -For text-to-image we use `stabilityai/stable-diffusion-2-1-unclip-small` as it was trained on CLIP ViT-L/14 embedding, the same as the Karlo model prior. [stabilityai/stable-diffusion-2-1-unclip](https://hf.co/stabilityai/stable-diffusion-2-1-unclip) was trained on OpenCLIP ViT-H, so we don't recommend its use. - - +> [!WARNING] +> For text-to-image we use `stabilityai/stable-diffusion-2-1-unclip-small` as it was trained on CLIP ViT-L/14 embedding, the same as the Karlo model prior. [stabilityai/stable-diffusion-2-1-unclip](https://hf.co/stabilityai/stable-diffusion-2-1-unclip) was trained on OpenCLIP ViT-H, so we don't recommend its use. ### Text guided Image-to-Image Variation @@ -99,11 +96,8 @@ image = pipe(init_image, prompt=prompt).images[0] image ``` - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## StableUnCLIPPipeline diff --git a/docs/source/en/api/pipelines/text_to_video.md b/docs/source/en/api/pipelines/text_to_video.md index 7faf88d133..d7c37be637 100644 --- a/docs/source/en/api/pipelines/text_to_video.md +++ b/docs/source/en/api/pipelines/text_to_video.md @@ -174,11 +174,8 @@ Video generation is memory-intensive and one way to reduce your memory usage is Check out the [Text or image-to-video](text-img2vid) guide for more details about how certain parameters can affect video generation and how to optimize inference by reducing memory usage. - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## TextToVideoSDPipeline [[autodoc]] TextToVideoSDPipeline diff --git a/docs/source/en/api/pipelines/text_to_video_zero.md b/docs/source/en/api/pipelines/text_to_video_zero.md index 5fe3789d82..50e7620760 100644 --- a/docs/source/en/api/pipelines/text_to_video_zero.md +++ b/docs/source/en/api/pipelines/text_to_video_zero.md @@ -289,11 +289,8 @@ can run with custom [DreamBooth](../../training/dreambooth) models, as shown bel You can filter out some available DreamBooth-trained models with [this link](https://huggingface.co/models?search=dreambooth). - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## TextToVideoZeroPipeline [[autodoc]] TextToVideoZeroPipeline diff --git a/docs/source/en/api/pipelines/unclip.md b/docs/source/en/api/pipelines/unclip.md index 8011a4b533..7c5c2b0d9a 100644 --- a/docs/source/en/api/pipelines/unclip.md +++ b/docs/source/en/api/pipelines/unclip.md @@ -20,11 +20,8 @@ The abstract from the paper is following: You can find lucidrains' DALL-E 2 recreation at [lucidrains/DALLE2-pytorch](https://github.com/lucidrains/DALLE2-pytorch). - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## UnCLIPPipeline [[autodoc]] UnCLIPPipeline diff --git a/docs/source/en/api/pipelines/unidiffuser.md b/docs/source/en/api/pipelines/unidiffuser.md index 7d767f2db5..2ff700e4b8 100644 --- a/docs/source/en/api/pipelines/unidiffuser.md +++ b/docs/source/en/api/pipelines/unidiffuser.md @@ -27,11 +27,8 @@ The abstract from the paper is: You can find the original codebase at [thu-ml/unidiffuser](https://github.com/thu-ml/unidiffuser) and additional checkpoints at [thu-ml](https://huggingface.co/thu-ml). - - -There is currently an issue on PyTorch 1.X where the output images are all black or the pixel values become `NaNs`. This issue can be mitigated by switching to PyTorch 2.X. - - +> [!WARNING] +> There is currently an issue on PyTorch 1.X where the output images are all black or the pixel values become `NaNs`. This issue can be mitigated by switching to PyTorch 2.X. This pipeline was contributed by [dg845](https://github.com/dg845). ❤️ @@ -197,11 +194,8 @@ final_prompt = sample.text[0] print(final_prompt) ``` - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## UniDiffuserPipeline [[autodoc]] UniDiffuserPipeline diff --git a/docs/source/en/api/pipelines/value_guided_sampling.md b/docs/source/en/api/pipelines/value_guided_sampling.md index 797847ee47..d050ea309c 100644 --- a/docs/source/en/api/pipelines/value_guided_sampling.md +++ b/docs/source/en/api/pipelines/value_guided_sampling.md @@ -12,11 +12,8 @@ specific language governing permissions and limitations under the License. # Value-guided planning - - -🧪 This is an experimental pipeline for reinforcement learning! - - +> [!WARNING] +> 🧪 This is an experimental pipeline for reinforcement learning! This pipeline is based on the [Planning with Diffusion for Flexible Behavior Synthesis](https://huggingface.co/papers/2205.09991) paper by Michael Janner, Yilun Du, Joshua B. Tenenbaum, Sergey Levine. @@ -28,11 +25,8 @@ You can find additional information about the model on the [project page](https: The script to run the model is available [here](https://github.com/huggingface/diffusers/tree/main/examples/reinforcement_learning). - - -Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - - +> [!TIP] +> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. ## ValueGuidedRLPipeline [[autodoc]] diffusers.experimental.ValueGuidedRLPipeline diff --git a/docs/source/en/api/quantization.md b/docs/source/en/api/quantization.md index 31271f1722..7fa7c7c9d0 100644 --- a/docs/source/en/api/quantization.md +++ b/docs/source/en/api/quantization.md @@ -15,11 +15,8 @@ specific language governing permissions and limitations under the License. Quantization techniques reduce memory and computational costs by representing weights and activations with lower-precision data types like 8-bit integers (int8). This enables loading larger models you normally wouldn't be able to fit into memory, and speeding up inference. - - -Learn how to quantize models in the [Quantization](../quantization/overview) guide. - - +> [!TIP] +> Learn how to quantize models in the [Quantization](../quantization/overview) guide. ## PipelineQuantizationConfig diff --git a/docs/source/en/api/schedulers/ddim.md b/docs/source/en/api/schedulers/ddim.md index 5d6b4673d2..61ef30c786 100644 --- a/docs/source/en/api/schedulers/ddim.md +++ b/docs/source/en/api/schedulers/ddim.md @@ -28,11 +28,8 @@ The original codebase of this paper can be found at [ermongroup/ddim](https://gi The paper [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) claims that a mismatch between the training and inference settings leads to suboptimal inference generation results for Stable Diffusion. To fix this, the authors propose: - - -🧪 This is an experimental feature! - - +> [!WARNING] +> 🧪 This is an experimental feature! 1. rescale the noise schedule to enforce zero terminal signal-to-noise ratio (SNR) diff --git a/docs/source/en/api/schedulers/score_sde_vp.md b/docs/source/en/api/schedulers/score_sde_vp.md index 0a1fe5a3be..8ce9800ee3 100644 --- a/docs/source/en/api/schedulers/score_sde_vp.md +++ b/docs/source/en/api/schedulers/score_sde_vp.md @@ -18,11 +18,8 @@ The abstract from the paper is: *Creating noise from data is easy; creating data from noise is generative modeling. We present a stochastic differential equation (SDE) that smoothly transforms a complex data distribution to a known prior distribution by slowly injecting noise, and a corresponding reverse-time SDE that transforms the prior distribution back into the data distribution by slowly removing the noise. Crucially, the reverse-time SDE depends only on the time-dependent gradient field (\aka, score) of the perturbed data distribution. By leveraging advances in score-based generative modeling, we can accurately estimate these scores with neural networks, and use numerical SDE solvers to generate samples. We show that this framework encapsulates previous approaches in score-based generative modeling and diffusion probabilistic modeling, allowing for new sampling procedures and new modeling capabilities. In particular, we introduce a predictor-corrector framework to correct errors in the evolution of the discretized reverse-time SDE. We also derive an equivalent neural ODE that samples from the same distribution as the SDE, but additionally enables exact likelihood computation, and improved sampling efficiency. In addition, we provide a new way to solve inverse problems with score-based models, as demonstrated with experiments on class-conditional generation, image inpainting, and colorization. Combined with multiple architectural improvements, we achieve record-breaking performance for unconditional image generation on CIFAR-10 with an Inception score of 9.89 and FID of 2.20, a competitive likelihood of 2.99 bits/dim, and demonstrate high fidelity generation of 1024 x 1024 images for the first time from a score-based generative model.* - - -🚧 This scheduler is under construction! - - +> [!WARNING] +> 🚧 This scheduler is under construction! ## ScoreSdeVpScheduler [[autodoc]] schedulers.deprecated.scheduling_sde_vp.ScoreSdeVpScheduler diff --git a/docs/source/en/conceptual/evaluation.md b/docs/source/en/conceptual/evaluation.md index 6546e5bf24..4af38254be 100644 --- a/docs/source/en/conceptual/evaluation.md +++ b/docs/source/en/conceptual/evaluation.md @@ -104,13 +104,10 @@ We can also set `num_images_per_prompt` accordingly to compare different images Once several images are generated from all the prompts using multiple models (under evaluation), these results are presented to human evaluators for scoring. For more details on the DrawBench and PartiPrompts benchmarks, refer to their respective papers. - - -It is useful to look at some inference samples while a model is training to measure the -training progress. In our [training scripts](https://github.com/huggingface/diffusers/tree/main/examples/), we support this utility with additional support for -logging to TensorBoard and Weights & Biases. - - +> [!TIP] +> It is useful to look at some inference samples while a model is training to measure the +> training progress. In our [training scripts](https://github.com/huggingface/diffusers/tree/main/examples/), we support this utility with additional support for +> logging to TensorBoard and Weights & Biases. ## Quantitative Evaluation @@ -205,14 +202,11 @@ print(f"CLIP Score with v-1-5: {sd_clip_score_1_5}") It seems like the [v1-5](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) checkpoint performs better than its predecessor. Note, however, that the number of prompts we used to compute the CLIP scores is quite low. For a more practical evaluation, this number should be way higher, and the prompts should be diverse. - - -By construction, there are some limitations in this score. The captions in the training dataset -were crawled from the web and extracted from `alt` and similar tags associated an image on the internet. -They are not necessarily representative of what a human being would use to describe an image. Hence we -had to "engineer" some prompts here. - - +> [!WARNING] +> By construction, there are some limitations in this score. The captions in the training dataset +> were crawled from the web and extracted from `alt` and similar tags associated an image on the internet. +> They are not necessarily representative of what a human being would use to describe an image. Hence we +> had to "engineer" some prompts here. ### Image-conditioned text-to-image generation @@ -421,11 +415,8 @@ We can extend the idea of this metric to measure how similar the original image We can use these metrics for similar pipelines such as the [`StableDiffusionPix2PixZeroPipeline`](https://huggingface.co/docs/diffusers/main/en/api/pipelines/pix2pix_zero#diffusers.StableDiffusionPix2PixZeroPipeline). - - -Both CLIP score and CLIP direction similarity rely on the CLIP model, which can make the evaluations biased. - - +> [!TIP] +> Both CLIP score and CLIP direction similarity rely on the CLIP model, which can make the evaluations biased. ***Extending metrics like IS, FID (discussed later), or KID can be difficult*** when the model under evaluation was pre-trained on a large image-captioning dataset (such as the [LAION-5B dataset](https://laion.ai/blog/laion-5b/)). This is because underlying these metrics is an InceptionNet (pre-trained on the ImageNet-1k dataset) used for extracting intermediate image features. The pre-training dataset of Stable Diffusion may have limited overlap with the pre-training dataset of InceptionNet, so it is not a good candidate here for feature extraction. @@ -554,21 +545,18 @@ The lower the FID, the better it is. Several things can influence FID here: For the last two points, it is, therefore, a good practice to run the evaluation across different seeds and inference steps, and then report an average result. - - -FID results tend to be fragile as they depend on a lot of factors: - -* The specific Inception model used during computation. -* The implementation accuracy of the computation. -* The image format (not the same if we start from PNGs vs JPGs). - -Keeping that in mind, FID is often most useful when comparing similar runs, but it is -hard to reproduce paper results unless the authors carefully disclose the FID -measurement code. - -These points apply to other related metrics too, such as KID and IS. - - +> [!WARNING] +> FID results tend to be fragile as they depend on a lot of factors: +> +> * The specific Inception model used during computation. +> * The implementation accuracy of the computation. +> * The image format (not the same if we start from PNGs vs JPGs). +> +> Keeping that in mind, FID is often most useful when comparing similar runs, but it is +> hard to reproduce paper results unless the authors carefully disclose the FID +> measurement code. +> +> These points apply to other related metrics too, such as KID and IS. As a final step, let's visually inspect the `fake_images`. diff --git a/docs/source/en/optimization/coreml.md b/docs/source/en/optimization/coreml.md index cd0e662bb7..71da1e3dc1 100644 --- a/docs/source/en/optimization/coreml.md +++ b/docs/source/en/optimization/coreml.md @@ -16,11 +16,8 @@ specific language governing permissions and limitations under the License. Core ML models can leverage all the compute engines available in Apple devices: the CPU, the GPU, and the Apple Neural Engine (or ANE, a tensor-optimized accelerator available in Apple Silicon Macs and modern iPhones/iPads). Depending on the model and the device it's running on, Core ML can mix and match compute engines too, so some portions of the model may run on the CPU while others run on GPU, for example. - - -You can also run the `diffusers` Python codebase on Apple Silicon Macs using the `mps` accelerator built into PyTorch. This approach is explained in depth in [the mps guide](mps), but it is not compatible with native apps. - - +> [!TIP] +> You can also run the `diffusers` Python codebase on Apple Silicon Macs using the `mps` accelerator built into PyTorch. This approach is explained in depth in [the mps guide](mps), but it is not compatible with native apps. ## Stable Diffusion Core ML Checkpoints diff --git a/docs/source/en/optimization/fp16.md b/docs/source/en/optimization/fp16.md index 76d749ecf3..941f53604c 100644 --- a/docs/source/en/optimization/fp16.md +++ b/docs/source/en/optimization/fp16.md @@ -239,11 +239,8 @@ The `step()` function is [called](https://github.com/huggingface/diffusers/blob/ In general, the `sigmas` should [stay on the CPU](https://github.com/huggingface/diffusers/blob/35a969d297cba69110d175ee79c59312b9f49e1e/src/diffusers/schedulers/scheduling_euler_discrete.py#L240) to avoid the communication sync and latency. - - -Refer to the [torch.compile and Diffusers: A Hands-On Guide to Peak Performance](https://pytorch.org/blog/torch-compile-and-diffusers-a-hands-on-guide-to-peak-performance/) blog post for maximizing performance with `torch.compile` for diffusion models. - - +> [!TIP] +> Refer to the [torch.compile and Diffusers: A Hands-On Guide to Peak Performance](https://pytorch.org/blog/torch-compile-and-diffusers-a-hands-on-guide-to-peak-performance/) blog post for maximizing performance with `torch.compile` for diffusion models. ### Benchmarks diff --git a/docs/source/en/optimization/mps.md b/docs/source/en/optimization/mps.md index 7e4c2716ac..b5afa25b2f 100644 --- a/docs/source/en/optimization/mps.md +++ b/docs/source/en/optimization/mps.md @@ -38,11 +38,8 @@ image = pipe(prompt).images[0] image ``` - - -The PyTorch [mps](https://pytorch.org/docs/stable/notes/mps.html) backend does not support NDArray sizes greater than `2**32`. Please open an [Issue](https://github.com/huggingface/diffusers/issues/new/choose) if you encounter this problem so we can investigate. - - +> [!WARNING] +> The PyTorch [mps](https://pytorch.org/docs/stable/notes/mps.html) backend does not support NDArray sizes greater than `2**32`. Please open an [Issue](https://github.com/huggingface/diffusers/issues/new/choose) if you encounter this problem so we can investigate. If you're using **PyTorch 1.13**, you need to "prime" the pipeline with an additional one-time pass through it. This is a temporary workaround for an issue where the first inference pass produces slightly different results than subsequent ones. You only need to do this pass once, and after just one inference step you can discard the result. diff --git a/docs/source/en/optimization/neuron.md b/docs/source/en/optimization/neuron.md index fa933317b4..6a45bd0563 100644 --- a/docs/source/en/optimization/neuron.md +++ b/docs/source/en/optimization/neuron.md @@ -20,11 +20,8 @@ Diffusers functionalities are available on [AWS Inf2 instances](https://aws.amaz python -m pip install --upgrade-strategy eager optimum[neuronx] ``` - - -We provide pre-built [Hugging Face Neuron Deep Learning AMI](https://aws.amazon.com/marketplace/pp/prodview-gr3e6yiscria2) (DLAMI) and Optimum Neuron containers for Amazon SageMaker. It's recommended to correctly set up your environment. - - +> [!TIP] +> We provide pre-built [Hugging Face Neuron Deep Learning AMI](https://aws.amazon.com/marketplace/pp/prodview-gr3e6yiscria2) (DLAMI) and Optimum Neuron containers for Amazon SageMaker. It's recommended to correctly set up your environment. The example below demonstrates how to generate images with the Stable Diffusion XL model on an inf2.8xlarge instance (you can switch to cheaper inf2.xlarge instances once the model is compiled). To generate some images, use the [`~optimum.neuron.NeuronStableDiffusionXLPipeline`] class, which is similar to the [`StableDiffusionXLPipeline`] class in Diffusers. diff --git a/docs/source/en/optimization/onnx.md b/docs/source/en/optimization/onnx.md index d160dcffe8..620f2af994 100644 --- a/docs/source/en/optimization/onnx.md +++ b/docs/source/en/optimization/onnx.md @@ -34,11 +34,8 @@ image = pipeline(prompt).images[0] pipeline.save_pretrained("./onnx-stable-diffusion-v1-5") ``` - - -Generating multiple prompts in a batch seems to take too much memory. While we look into it, you may need to iterate instead of batching. - - +> [!WARNING] +> Generating multiple prompts in a batch seems to take too much memory. While we look into it, you may need to iterate instead of batching. To export the pipeline in the ONNX format offline and use it later for inference, use the [`optimum-cli export`](https://huggingface.co/docs/optimum/main/en/exporters/onnx/usage_guides/export_a_model#exporting-a-model-to-onnx-using-the-cli) command: diff --git a/docs/source/en/optimization/xformers.md b/docs/source/en/optimization/xformers.md index 3e2792fd5f..523e815595 100644 --- a/docs/source/en/optimization/xformers.md +++ b/docs/source/en/optimization/xformers.md @@ -20,16 +20,10 @@ Install xFormers from `pip`: pip install xformers ``` - - -The xFormers `pip` package requires the latest version of PyTorch. If you need to use a previous version of PyTorch, then we recommend [installing xFormers from the source](https://github.com/facebookresearch/xformers#installing-xformers). - - +> [!TIP] +> The xFormers `pip` package requires the latest version of PyTorch. If you need to use a previous version of PyTorch, then we recommend [installing xFormers from the source](https://github.com/facebookresearch/xformers#installing-xformers). After xFormers is installed, you can use `enable_xformers_memory_efficient_attention()` for faster inference and reduced memory consumption as shown in this [section](memory#memory-efficient-attention). - - -According to this [issue](https://github.com/huggingface/diffusers/issues/2234#issuecomment-1416931212), xFormers `v0.0.16` cannot be used for training (fine-tune or DreamBooth) in some GPUs. If you observe this problem, please install a development version as indicated in the issue comments. - - +> [!WARNING] +> According to this [issue](https://github.com/huggingface/diffusers/issues/2234#issuecomment-1416931212), xFormers `v0.0.16` cannot be used for training (fine-tune or DreamBooth) in some GPUs. If you observe this problem, please install a development version as indicated in the issue comments. diff --git a/docs/source/en/quantization/bitsandbytes.md b/docs/source/en/quantization/bitsandbytes.md index f97119d5f4..0729472744 100644 --- a/docs/source/en/quantization/bitsandbytes.md +++ b/docs/source/en/quantization/bitsandbytes.md @@ -206,11 +206,8 @@ Once a model is quantized, you can push the model to the Hub with the [`~ModelMi - - -Training with 8-bit and 4-bit weights are only supported for training *extra* parameters. - - +> [!WARNING] +> Training with 8-bit and 4-bit weights are only supported for training *extra* parameters. Check your memory footprint with the `get_memory_footprint` method: @@ -234,11 +231,8 @@ model_4bit = AutoModel.from_pretrained( ## 8-bit (LLM.int8() algorithm) - - -Learn more about the details of 8-bit quantization in this [blog post](https://huggingface.co/blog/hf-bitsandbytes-integration)! - - +> [!TIP] +> Learn more about the details of 8-bit quantization in this [blog post](https://huggingface.co/blog/hf-bitsandbytes-integration)! This section explores some of the specific features of 8-bit models, such as outlier thresholds and skipping module conversion. @@ -283,11 +277,8 @@ model_8bit = SD3Transformer2DModel.from_pretrained( ## 4-bit (QLoRA algorithm) - - -Learn more about its details in this [blog post](https://huggingface.co/blog/4bit-transformers-bitsandbytes). - - +> [!TIP] +> Learn more about its details in this [blog post](https://huggingface.co/blog/4bit-transformers-bitsandbytes). This section explores some of the specific features of 4-bit models, such as changing the compute data type, using the Normal Float 4 (NF4) data type, and using nested quantization. diff --git a/docs/source/en/training/controlnet.md b/docs/source/en/training/controlnet.md index 17da819db8..840130d2b4 100644 --- a/docs/source/en/training/controlnet.md +++ b/docs/source/en/training/controlnet.md @@ -33,11 +33,8 @@ cd examples/controlnet pip install -r requirements.txt ``` - - -🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It'll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate [Quick tour](https://huggingface.co/docs/accelerate/quicktour) to learn more. - - +> [!TIP] +> 🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It'll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate [Quick tour](https://huggingface.co/docs/accelerate/quicktour) to learn more. Initialize an 🤗 Accelerate environment: @@ -61,11 +58,8 @@ write_basic_config() Lastly, if you want to train a model on your own dataset, take a look at the [Create a dataset for training](create_dataset) guide to learn how to create a dataset that works with the training script. - - -The following sections highlight parts of the training script that are important for understanding how to modify it, but it doesn't cover every aspect of the script in detail. If you're interested in learning more, feel free to read through the [script](https://github.com/huggingface/diffusers/blob/main/examples/controlnet/train_controlnet.py) and let us know if you have any questions or concerns. - - +> [!TIP] +> The following sections highlight parts of the training script that are important for understanding how to modify it, but it doesn't cover every aspect of the script in detail. If you're interested in learning more, feel free to read through the [script](https://github.com/huggingface/diffusers/blob/main/examples/controlnet/train_controlnet.py) and let us know if you have any questions or concerns. ## Script parameters @@ -100,11 +94,8 @@ As with the script parameters, a general walkthrough of the training script is p The training script has a [`make_train_dataset`](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/controlnet/train_controlnet.py#L582) function for preprocessing the dataset with image transforms and caption tokenization. You'll see that in addition to the usual caption tokenization and image transforms, the script also includes transforms for the conditioning image. - - -If you're streaming a dataset on a TPU, performance may be bottlenecked by the 🤗 Datasets library which is not optimized for images. To ensure maximum throughput, you're encouraged to explore other dataset formats like [WebDataset](https://webdataset.github.io/webdataset/), [TorchData](https://github.com/pytorch/data), and [TensorFlow Datasets](https://www.tensorflow.org/datasets/tfless_tfds). - - +> [!TIP] +> If you're streaming a dataset on a TPU, performance may be bottlenecked by the 🤗 Datasets library which is not optimized for images. To ensure maximum throughput, you're encouraged to explore other dataset formats like [WebDataset](https://webdataset.github.io/webdataset/), [TorchData](https://github.com/pytorch/data), and [TensorFlow Datasets](https://www.tensorflow.org/datasets/tfless_tfds). ```py conditioning_image_transforms = transforms.Compose( diff --git a/docs/source/en/training/create_dataset.md b/docs/source/en/training/create_dataset.md index 8e0d6f9200..725f143bba 100644 --- a/docs/source/en/training/create_dataset.md +++ b/docs/source/en/training/create_dataset.md @@ -7,11 +7,8 @@ This guide will show you two ways to create a dataset to finetune on: - provide a folder of images to the `--train_data_dir` argument - upload a dataset to the Hub and pass the dataset repository id to the `--dataset_name` argument - - -💡 Learn more about how to create an image dataset for training in the [Create an image dataset](https://huggingface.co/docs/datasets/image_dataset) guide. - - +> [!TIP] +> 💡 Learn more about how to create an image dataset for training in the [Create an image dataset](https://huggingface.co/docs/datasets/image_dataset) guide. ## Provide a dataset as a folder @@ -33,11 +30,8 @@ accelerate launch train_unconditional.py \ ## Upload your data to the Hub - - -💡 For more details and context about creating and uploading a dataset to the Hub, take a look at the [Image search with 🤗 Datasets](https://huggingface.co/blog/image-search-datasets) post. - - +> [!TIP] +> 💡 For more details and context about creating and uploading a dataset to the Hub, take a look at the [Image search with 🤗 Datasets](https://huggingface.co/blog/image-search-datasets) post. Start by creating a dataset with the [`ImageFolder`](https://huggingface.co/docs/datasets/image_load#imagefolder) feature, which creates an `image` column containing the PIL-encoded images. diff --git a/docs/source/en/training/custom_diffusion.md b/docs/source/en/training/custom_diffusion.md index e803448b5f..bfa4fe6f9e 100644 --- a/docs/source/en/training/custom_diffusion.md +++ b/docs/source/en/training/custom_diffusion.md @@ -34,11 +34,8 @@ pip install -r requirements.txt pip install clip-retrieval ``` - - -🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It'll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate [Quick tour](https://huggingface.co/docs/accelerate/quicktour) to learn more. - - +> [!TIP] +> 🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It'll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate [Quick tour](https://huggingface.co/docs/accelerate/quicktour) to learn more. Initialize an 🤗 Accelerate environment: @@ -62,11 +59,8 @@ write_basic_config() Lastly, if you want to train a model on your own dataset, take a look at the [Create a dataset for training](create_dataset) guide to learn how to create a dataset that works with the training script. - - -The following sections highlight parts of the training script that are important for understanding how to modify it, but it doesn't cover every aspect of the script in detail. If you're interested in learning more, feel free to read through the [script](https://github.com/huggingface/diffusers/blob/main/examples/custom_diffusion/train_custom_diffusion.py) and let us know if you have any questions or concerns. - - +> [!TIP] +> The following sections highlight parts of the training script that are important for understanding how to modify it, but it doesn't cover every aspect of the script in detail. If you're interested in learning more, feel free to read through the [script](https://github.com/huggingface/diffusers/blob/main/examples/custom_diffusion/train_custom_diffusion.py) and let us know if you have any questions or concerns. ## Script parameters @@ -117,11 +111,8 @@ accelerate launch train_custom_diffusion.py \ ## Training script - - -A lot of the code in the Custom Diffusion training script is similar to the [DreamBooth](dreambooth#training-script) script. This guide instead focuses on the code that is relevant to Custom Diffusion. - - +> [!TIP] +> A lot of the code in the Custom Diffusion training script is similar to the [DreamBooth](dreambooth#training-script) script. This guide instead focuses on the code that is relevant to Custom Diffusion. The Custom Diffusion training script has two dataset classes: @@ -224,16 +215,13 @@ Set the environment variable `MODEL_NAME` to a model id on the Hub or a path to To monitor training progress with Weights and Biases, add the `--report_to=wandb` parameter to the training command and specify a validation prompt with `--validation_prompt`. This is useful for debugging and saving intermediate results. - - -If you're training on human faces, the Custom Diffusion team has found the following parameters to work well: - -- `--learning_rate=5e-6` -- `--max_train_steps` can be anywhere between 1000 and 2000 -- `--freeze_model=crossattn` -- use at least 15-20 images to train with - - +> [!TIP] +> If you're training on human faces, the Custom Diffusion team has found the following parameters to work well: +> +> - `--learning_rate=5e-6` +> - `--max_train_steps` can be anywhere between 1000 and 2000 +> - `--freeze_model=crossattn` +> - use at least 15-20 images to train with diff --git a/docs/source/en/training/distributed_inference.md b/docs/source/en/training/distributed_inference.md index 586f765709..f9756e1a67 100644 --- a/docs/source/en/training/distributed_inference.md +++ b/docs/source/en/training/distributed_inference.md @@ -56,6 +56,9 @@ Call `accelerate launch` to run the script and use the `--num_processes` argumen accelerate launch run_distributed.py --num_processes=2 ``` +> [!TIP] +> Refer to this minimal example [script](https://gist.github.com/sayakpaul/cfaebd221820d7b43fae638b4dfa01ba) for running inference across multiple GPUs. To learn more, take a look at the [Distributed Inference with 🤗 Accelerate](https://huggingface.co/docs/accelerate/en/usage_guides/distributed_inference#distributed-inference-with-accelerate) guide. + ## PyTorch Distributed PyTorch [DistributedDataParallel](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html) enables [data parallelism](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=data_parallelism), which replicates the same model on each device, to process different batches of data in parallel. diff --git a/docs/source/en/training/dreambooth.md b/docs/source/en/training/dreambooth.md index 3a5ba5aa39..81ed09c9d0 100644 --- a/docs/source/en/training/dreambooth.md +++ b/docs/source/en/training/dreambooth.md @@ -33,11 +33,8 @@ cd examples/dreambooth pip install -r requirements.txt ``` - - -🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It'll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate [Quick tour](https://huggingface.co/docs/accelerate/quicktour) to learn more. - - +> [!TIP] +> 🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It'll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate [Quick tour](https://huggingface.co/docs/accelerate/quicktour) to learn more. Initialize an 🤗 Accelerate environment: @@ -61,19 +58,13 @@ write_basic_config() Lastly, if you want to train a model on your own dataset, take a look at the [Create a dataset for training](create_dataset) guide to learn how to create a dataset that works with the training script. - - -The following sections highlight parts of the training script that are important for understanding how to modify it, but it doesn't cover every aspect of the script in detail. If you're interested in learning more, feel free to read through the [script](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth.py) and let us know if you have any questions or concerns. - - +> [!TIP] +> The following sections highlight parts of the training script that are important for understanding how to modify it, but it doesn't cover every aspect of the script in detail. If you're interested in learning more, feel free to read through the [script](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth.py) and let us know if you have any questions or concerns. ## Script parameters - - -DreamBooth is very sensitive to training hyperparameters, and it is easy to overfit. Read the [Training Stable Diffusion with Dreambooth using 🧨 Diffusers](https://huggingface.co/blog/dreambooth) blog post for recommended settings for different subjects to help you choose the appropriate hyperparameters. - - +> [!WARNING] +> DreamBooth is very sensitive to training hyperparameters, and it is easy to overfit. Read the [Training Stable Diffusion with Dreambooth using 🧨 Diffusers](https://huggingface.co/blog/dreambooth) blog post for recommended settings for different subjects to help you choose the appropriate hyperparameters. The training script offers many parameters for customizing your training run. All of the parameters and their descriptions are found in the [`parse_args()`](https://github.com/huggingface/diffusers/blob/072e00897a7cf4302c347a63ec917b4b8add16d4/examples/dreambooth/train_dreambooth.py#L228) function. The parameters are set with default values that should work pretty well out-of-the-box, but you can also set your own values in the training command if you'd like. @@ -319,29 +310,26 @@ accelerate launch train_dreambooth.py \ Once training is complete, you can use your newly trained model for inference! - - -Can't wait to try your model for inference before training is complete? 🤭 Make sure you have the latest version of 🤗 Accelerate installed. - -```py -from diffusers import DiffusionPipeline, UNet2DConditionModel -from transformers import CLIPTextModel -import torch - -unet = UNet2DConditionModel.from_pretrained("path/to/model/checkpoint-100/unet") - -# if you have trained with `--args.train_text_encoder` make sure to also load the text encoder -text_encoder = CLIPTextModel.from_pretrained("path/to/model/checkpoint-100/checkpoint-100/text_encoder") - -pipeline = DiffusionPipeline.from_pretrained( - "stable-diffusion-v1-5/stable-diffusion-v1-5", unet=unet, text_encoder=text_encoder, dtype=torch.float16, -).to("cuda") - -image = pipeline("A photo of sks dog in a bucket", num_inference_steps=50, guidance_scale=7.5).images[0] -image.save("dog-bucket.png") -``` - - +> [!TIP] +> Can't wait to try your model for inference before training is complete? 🤭 Make sure you have the latest version of 🤗 Accelerate installed. +> +> ```py +> from diffusers import DiffusionPipeline, UNet2DConditionModel +> from transformers import CLIPTextModel +> import torch +> +> unet = UNet2DConditionModel.from_pretrained("path/to/model/checkpoint-100/unet") +> +> # if you have trained with `--args.train_text_encoder` make sure to also load the text encoder +> text_encoder = CLIPTextModel.from_pretrained("path/to/model/checkpoint-100/checkpoint-100/text_encoder") +> +> pipeline = DiffusionPipeline.from_pretrained( +> "stable-diffusion-v1-5/stable-diffusion-v1-5", unet=unet, text_encoder=text_encoder, dtype=torch.float16, +> ).to("cuda") +> +> image = pipeline("A photo of sks dog in a bucket", num_inference_steps=50, guidance_scale=7.5).images[0] +> image.save("dog-bucket.png") +> ``` ```py from diffusers import DiffusionPipeline diff --git a/docs/source/en/training/instructpix2pix.md b/docs/source/en/training/instructpix2pix.md index c1ba5d870a..a1c94bb33f 100644 --- a/docs/source/en/training/instructpix2pix.md +++ b/docs/source/en/training/instructpix2pix.md @@ -31,11 +31,8 @@ cd examples/instruct_pix2pix pip install -r requirements.txt ``` - - -🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It'll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate [Quick tour](https://huggingface.co/docs/accelerate/quicktour) to learn more. - - +> [!TIP] +> 🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It'll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate [Quick tour](https://huggingface.co/docs/accelerate/quicktour) to learn more. Initialize an 🤗 Accelerate environment: @@ -59,11 +56,8 @@ write_basic_config() Lastly, if you want to train a model on your own dataset, take a look at the [Create a dataset for training](create_dataset) guide to learn how to create a dataset that works with the training script. - - -The following sections highlight parts of the training script that are important for understanding how to modify it, but it doesn't cover every aspect of the script in detail. If you're interested in learning more, feel free to read through the [script](https://github.com/huggingface/diffusers/blob/main/examples/instruct_pix2pix/train_instruct_pix2pix.py) and let us know if you have any questions or concerns. - - +> [!TIP] +> The following sections highlight parts of the training script that are important for understanding how to modify it, but it doesn't cover every aspect of the script in detail. If you're interested in learning more, feel free to read through the [script](https://github.com/huggingface/diffusers/blob/main/examples/instruct_pix2pix/train_instruct_pix2pix.py) and let us know if you have any questions or concerns. ## Script parameters @@ -174,15 +168,12 @@ This guide uses the [fusing/instructpix2pix-1000-samples](https://huggingface.co Set the `MODEL_NAME` environment variable to the name of the model (can be a model id on the Hub or a path to a local model), and the `DATASET_ID` to the name of the dataset on the Hub. The script creates and saves all the components (feature extractor, scheduler, text encoder, UNet, etc.) to a subfolder in your repository. - - -For better results, try longer training runs with a larger dataset. We've only tested this training script on a smaller-scale dataset. - -
- -To monitor training progress with Weights and Biases, add the `--report_to=wandb` parameter to the training command and specify a validation image with `--val_image_url` and a validation prompt with `--validation_prompt`. This can be really useful for debugging the model. - - +> [!TIP] +> For better results, try longer training runs with a larger dataset. We've only tested this training script on a smaller-scale dataset. +> +>
+> +> To monitor training progress with Weights and Biases, add the `--report_to=wandb` parameter to the training command and specify a validation image with `--val_image_url` and a validation prompt with `--validation_prompt`. This can be really useful for debugging the model. If you’re training on more than one GPU, add the `--multi_gpu` parameter to the `accelerate launch` command. diff --git a/docs/source/en/training/kandinsky.md b/docs/source/en/training/kandinsky.md index 561bc1c351..6cfd9f8d60 100644 --- a/docs/source/en/training/kandinsky.md +++ b/docs/source/en/training/kandinsky.md @@ -12,11 +12,8 @@ specific language governing permissions and limitations under the License. # Kandinsky 2.2 - - -This script is experimental, and it's easy to overfit and run into issues like catastrophic forgetting. Try exploring different hyperparameters to get the best results on your dataset. - - +> [!WARNING] +> This script is experimental, and it's easy to overfit and run into issues like catastrophic forgetting. Try exploring different hyperparameters to get the best results on your dataset. Kandinsky 2.2 is a multilingual text-to-image model capable of producing more photorealistic images. The model includes an image prior model for creating image embeddings from text prompts, and a decoder model that generates images based on the prior model's embeddings. That's why you'll find two separate scripts in Diffusers for Kandinsky 2.2, one for training the prior model and one for training the decoder model. You can train both models separately, but to get the best results, you should train both the prior and decoder models. @@ -39,11 +36,8 @@ cd examples/kandinsky2_2/text_to_image pip install -r requirements.txt ``` - - -🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It'll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate [Quick tour](https://huggingface.co/docs/accelerate/quicktour) to learn more. - - +> [!TIP] +> 🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It'll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate [Quick tour](https://huggingface.co/docs/accelerate/quicktour) to learn more. Initialize an 🤗 Accelerate environment: @@ -67,11 +61,8 @@ write_basic_config() Lastly, if you want to train a model on your own dataset, take a look at the [Create a dataset for training](create_dataset) guide to learn how to create a dataset that works with the training script. - - -The following sections highlight parts of the training scripts that are important for understanding how to modify it, but it doesn't cover every aspect of the scripts in detail. If you're interested in learning more, feel free to read through the scripts and let us know if you have any questions or concerns. - - +> [!TIP] +> The following sections highlight parts of the training scripts that are important for understanding how to modify it, but it doesn't cover every aspect of the scripts in detail. If you're interested in learning more, feel free to read through the scripts and let us know if you have any questions or concerns. ## Script parameters @@ -209,11 +200,8 @@ You'll train on the [Naruto BLIP captions](https://huggingface.co/datasets/lambd If you’re training on more than one GPU, add the `--multi_gpu` parameter to the `accelerate launch` command. - - -To monitor training progress with Weights & Biases, add the `--report_to=wandb` parameter to the training command. You’ll also need to add the `--validation_prompt` to the training command to keep track of results. This can be really useful for debugging the model and viewing intermediate results. - - +> [!TIP] +> To monitor training progress with Weights & Biases, add the `--report_to=wandb` parameter to the training command. You’ll also need to add the `--validation_prompt` to the training command to keep track of results. This can be really useful for debugging the model and viewing intermediate results. @@ -283,11 +271,8 @@ prompt="A robot naruto, 4k photo" image = pipeline(prompt=prompt, negative_prompt=negative_prompt).images[0] ``` - - -Feel free to replace `kandinsky-community/kandinsky-2-2-decoder` with your own trained decoder checkpoint! - - +> [!TIP] +> Feel free to replace `kandinsky-community/kandinsky-2-2-decoder` with your own trained decoder checkpoint! diff --git a/docs/source/en/training/lcm_distill.md b/docs/source/en/training/lcm_distill.md index 280b6469f6..232f2eceed 100644 --- a/docs/source/en/training/lcm_distill.md +++ b/docs/source/en/training/lcm_distill.md @@ -33,11 +33,8 @@ cd examples/consistency_distillation pip install -r requirements.txt ``` - - -🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It'll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate [Quick tour](https://huggingface.co/docs/accelerate/quicktour) to learn more. - - +> [!TIP] +> 🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It'll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate [Quick tour](https://huggingface.co/docs/accelerate/quicktour) to learn more. Initialize an 🤗 Accelerate environment (try enabling `torch.compile` to significantly speedup training): @@ -63,11 +60,8 @@ Lastly, if you want to train a model on your own dataset, take a look at the [Cr ## Script parameters - - -The following sections highlight parts of the training script that are important for understanding how to modify it, but it doesn't cover every aspect of the script in detail. If you're interested in learning more, feel free to read through the [script](https://github.com/huggingface/diffusers/blob/main/examples/consistency_distillation/train_lcm_distill_sd_wds.py) and let us know if you have any questions or concerns. - - +> [!TIP] +> The following sections highlight parts of the training script that are important for understanding how to modify it, but it doesn't cover every aspect of the script in detail. If you're interested in learning more, feel free to read through the [script](https://github.com/huggingface/diffusers/blob/main/examples/consistency_distillation/train_lcm_distill_sd_wds.py) and let us know if you have any questions or concerns. The training script provides many parameters to help you customize your training run. All of the parameters and their descriptions are found in the [`parse_args()`](https://github.com/huggingface/diffusers/blob/3b37488fa3280aed6a95de044d7a42ffdcb565ef/examples/consistency_distillation/train_lcm_distill_sd_wds.py#L419) function. This function provides default values for each parameter, such as the training batch size and learning rate, but you can also set your own values in the training command if you'd like. diff --git a/docs/source/en/training/lora.md b/docs/source/en/training/lora.md index e97d8acdac..45a234b76a 100644 --- a/docs/source/en/training/lora.md +++ b/docs/source/en/training/lora.md @@ -12,19 +12,13 @@ specific language governing permissions and limitations under the License. # LoRA - - -This is experimental and the API may change in the future. - - +> [!WARNING] +> This is experimental and the API may change in the future. [LoRA (Low-Rank Adaptation of Large Language Models)](https://hf.co/papers/2106.09685) is a popular and lightweight training technique that significantly reduces the number of trainable parameters. It works by inserting a smaller number of new weights into the model and only these are trained. This makes training with LoRA much faster, memory-efficient, and produces smaller model weights (a few hundred MBs), which are easier to store and share. LoRA can also be combined with other training techniques like DreamBooth to speedup training. - - -LoRA is very versatile and supported for [DreamBooth](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth_lora.py), [Kandinsky 2.2](https://github.com/huggingface/diffusers/blob/main/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_decoder.py), [Stable Diffusion XL](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_lora_sdxl.py), [text-to-image](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_lora.py), and [Wuerstchen](https://github.com/huggingface/diffusers/blob/main/examples/wuerstchen/text_to_image/train_text_to_image_lora_prior.py). - - +> [!TIP] +> LoRA is very versatile and supported for [DreamBooth](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth_lora.py), [Kandinsky 2.2](https://github.com/huggingface/diffusers/blob/main/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_decoder.py), [Stable Diffusion XL](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_lora_sdxl.py), [text-to-image](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_lora.py), and [Wuerstchen](https://github.com/huggingface/diffusers/blob/main/examples/wuerstchen/text_to_image/train_text_to_image_lora_prior.py). This guide will explore the [train_text_to_image_lora.py](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_lora.py) script to help you become more familiar with it, and how you can adapt it for your own use-case. @@ -43,11 +37,8 @@ cd examples/text_to_image pip install -r requirements.txt ``` - - -🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It'll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate [Quick tour](https://huggingface.co/docs/accelerate/quicktour) to learn more. - - +> [!TIP] +> 🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It'll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate [Quick tour](https://huggingface.co/docs/accelerate/quicktour) to learn more. Initialize an 🤗 Accelerate environment: @@ -71,11 +62,8 @@ write_basic_config() Lastly, if you want to train a model on your own dataset, take a look at the [Create a dataset for training](create_dataset) guide to learn how to create a dataset that works with the training script. - - -The following sections highlight parts of the training script that are important for understanding how to modify it, but it doesn't cover every aspect of the script in detail. If you're interested in learning more, feel free to read through the [script](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_lora.py) and let us know if you have any questions or concerns. - - +> [!TIP] +> The following sections highlight parts of the training script that are important for understanding how to modify it, but it doesn't cover every aspect of the script in detail. If you're interested in learning more, feel free to read through the [script](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_lora.py) and let us know if you have any questions or concerns. ## Script parameters @@ -163,11 +151,8 @@ Let's train on the [Naruto BLIP captions](https://huggingface.co/datasets/lambda If you're training on more than one GPU, add the `--multi_gpu` parameter to the `accelerate launch` command. - - -A full training run takes ~5 hours on a 2080 Ti GPU with 11GB of VRAM. - - +> [!WARNING] +> A full training run takes ~5 hours on a 2080 Ti GPU with 11GB of VRAM. ```bash export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5" diff --git a/docs/source/en/training/sdxl.md b/docs/source/en/training/sdxl.md index 12051b7c2d..266bbc7d61 100644 --- a/docs/source/en/training/sdxl.md +++ b/docs/source/en/training/sdxl.md @@ -12,11 +12,8 @@ specific language governing permissions and limitations under the License. # Stable Diffusion XL - - -This script is experimental, and it's easy to overfit and run into issues like catastrophic forgetting. Try exploring different hyperparameters to get the best results on your dataset. - - +> [!WARNING] +> This script is experimental, and it's easy to overfit and run into issues like catastrophic forgetting. Try exploring different hyperparameters to get the best results on your dataset. [Stable Diffusion XL (SDXL)](https://hf.co/papers/2307.01952) is a larger and more powerful iteration of the Stable Diffusion model, capable of producing higher resolution images. @@ -39,11 +36,8 @@ cd examples/text_to_image pip install -r requirements_sdxl.txt ``` - - -🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It'll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate [Quick tour](https://huggingface.co/docs/accelerate/quicktour) to learn more. - - +> [!TIP] +> 🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It'll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate [Quick tour](https://huggingface.co/docs/accelerate/quicktour) to learn more. Initialize an 🤗 Accelerate environment: @@ -69,11 +63,8 @@ Lastly, if you want to train a model on your own dataset, take a look at the [Cr ## Script parameters - - -The following sections highlight parts of the training script that are important for understanding how to modify it, but it doesn't cover every aspect of the script in detail. If you're interested in learning more, feel free to read through the [script](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_sdxl.py) and let us know if you have any questions or concerns. - - +> [!TIP] +> The following sections highlight parts of the training script that are important for understanding how to modify it, but it doesn't cover every aspect of the script in detail. If you're interested in learning more, feel free to read through the [script](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_sdxl.py) and let us know if you have any questions or concerns. The training script provides many parameters to help you customize your training run. All of the parameters and their descriptions are found in the [`parse_args()`](https://github.com/huggingface/diffusers/blob/aab6de22c33cc01fb7bc81c0807d6109e2c998c9/examples/text_to_image/train_text_to_image_sdxl.py#L129) function. This function provides default values for each parameter, such as the training batch size and learning rate, but you can also set your own values in the training command if you'd like. @@ -178,11 +169,8 @@ Once you’ve made all your changes or you’re okay with the default configurat Let’s train on the [Naruto BLIP captions](https://huggingface.co/datasets/lambdalabs/naruto-blip-captions) dataset to generate your own Naruto characters. Set the environment variables `MODEL_NAME` and `DATASET_NAME` to the model and the dataset (either from the Hub or a local path). You should also specify a VAE other than the SDXL VAE (either from the Hub or a local path) with `VAE_NAME` to avoid numerical instabilities. - - -To monitor training progress with Weights & Biases, add the `--report_to=wandb` parameter to the training command. You’ll also need to add the `--validation_prompt` and `--validation_epochs` to the training command to keep track of results. This can be really useful for debugging the model and viewing intermediate results. - - +> [!TIP] +> To monitor training progress with Weights & Biases, add the `--report_to=wandb` parameter to the training command. You’ll also need to add the `--validation_prompt` and `--validation_epochs` to the training command to keep track of results. This can be really useful for debugging the model and viewing intermediate results. ```bash export MODEL_NAME="stabilityai/stable-diffusion-xl-base-1.0" diff --git a/docs/source/en/training/t2i_adapters.md b/docs/source/en/training/t2i_adapters.md index 243c591bea..6d76004073 100644 --- a/docs/source/en/training/t2i_adapters.md +++ b/docs/source/en/training/t2i_adapters.md @@ -33,11 +33,8 @@ cd examples/t2i_adapter pip install -r requirements.txt ``` - - -🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It'll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate [Quick tour](https://huggingface.co/docs/accelerate/quicktour) to learn more. - - +> [!TIP] +> 🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It'll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate [Quick tour](https://huggingface.co/docs/accelerate/quicktour) to learn more. Initialize an 🤗 Accelerate environment: @@ -61,11 +58,8 @@ write_basic_config() Lastly, if you want to train a model on your own dataset, take a look at the [Create a dataset for training](create_dataset) guide to learn how to create a dataset that works with the training script. - - -The following sections highlight parts of the training script that are important for understanding how to modify it, but it doesn't cover every aspect of the script in detail. If you're interested in learning more, feel free to read through the [script](https://github.com/huggingface/diffusers/blob/main/examples/t2i_adapter/train_t2i_adapter_sdxl.py) and let us know if you have any questions or concerns. - - +> [!TIP] +> The following sections highlight parts of the training script that are important for understanding how to modify it, but it doesn't cover every aspect of the script in detail. If you're interested in learning more, feel free to read through the [script](https://github.com/huggingface/diffusers/blob/main/examples/t2i_adapter/train_t2i_adapter_sdxl.py) and let us know if you have any questions or concerns. ## Script parameters @@ -166,11 +160,8 @@ wget https://huggingface.co/datasets/huggingface/documentation-images/resolve/ma wget https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_training/conditioning_image_2.png ``` - - -To monitor training progress with Weights & Biases, add the `--report_to=wandb` parameter to the training command. You'll also need to add the `--validation_image`, `--validation_prompt`, and `--validation_steps` to the training command to keep track of results. This can be really useful for debugging the model and viewing intermediate results. - - +> [!TIP] +> To monitor training progress with Weights & Biases, add the `--report_to=wandb` parameter to the training command. You'll also need to add the `--validation_image`, `--validation_prompt`, and `--validation_steps` to the training command to keep track of results. This can be really useful for debugging the model and viewing intermediate results. ```bash export MODEL_DIR="stabilityai/stable-diffusion-xl-base-1.0" diff --git a/docs/source/en/training/text2image.md b/docs/source/en/training/text2image.md index 5212fe8393..a9327457c7 100644 --- a/docs/source/en/training/text2image.md +++ b/docs/source/en/training/text2image.md @@ -12,11 +12,8 @@ specific language governing permissions and limitations under the License. # Text-to-image - - -The text-to-image script is experimental, and it's easy to overfit and run into issues like catastrophic forgetting. Try exploring different hyperparameters to get the best results on your dataset. - - +> [!WARNING] +> The text-to-image script is experimental, and it's easy to overfit and run into issues like catastrophic forgetting. Try exploring different hyperparameters to get the best results on your dataset. Text-to-image models like Stable Diffusion are conditioned to generate images given a text prompt. @@ -39,11 +36,8 @@ cd examples/text_to_image pip install -r requirements.txt ``` - - -🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It'll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate [Quick tour](https://huggingface.co/docs/accelerate/quicktour) to learn more. - - +> [!TIP] +> 🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It'll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate [Quick tour](https://huggingface.co/docs/accelerate/quicktour) to learn more. Initialize an 🤗 Accelerate environment: @@ -69,11 +63,8 @@ Lastly, if you want to train a model on your own dataset, take a look at the [Cr ## Script parameters - - -The following sections highlight parts of the training script that are important for understanding how to modify it, but it doesn't cover every aspect of the script in detail. If you're interested in learning more, feel free to read through the [script](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py) and let us know if you have any questions or concerns. - - +> [!TIP] +> The following sections highlight parts of the training script that are important for understanding how to modify it, but it doesn't cover every aspect of the script in detail. If you're interested in learning more, feel free to read through the [script](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py) and let us know if you have any questions or concerns. The training script provides many parameters to help you customize your training run. All of the parameters and their descriptions are found in the [`parse_args()`](https://github.com/huggingface/diffusers/blob/8959c5b9dec1c94d6ba482c94a58d2215c5fd026/examples/text_to_image/train_text_to_image.py#L193) function. This function provides default values for each parameter, such as the training batch size and learning rate, but you can also set your own values in the training command if you'd like. @@ -147,11 +138,8 @@ Once you've made all your changes or you're okay with the default configuration, Let's train on the [Naruto BLIP captions](https://huggingface.co/datasets/lambdalabs/naruto-blip-captions) dataset to generate your own Naruto characters. Set the environment variables `MODEL_NAME` and `dataset_name` to the model and the dataset (either from the Hub or a local path). If you're training on more than one GPU, add the `--multi_gpu` parameter to the `accelerate launch` command. - - -To train on a local dataset, set the `TRAIN_DIR` and `OUTPUT_DIR` environment variables to the path of the dataset and where to save the model to. - - +> [!TIP] +> To train on a local dataset, set the `TRAIN_DIR` and `OUTPUT_DIR` environment variables to the path of the dataset and where to save the model to. ```bash export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5" diff --git a/docs/source/en/training/text_inversion.md b/docs/source/en/training/text_inversion.md index 91af2f6afb..0b540107e9 100644 --- a/docs/source/en/training/text_inversion.md +++ b/docs/source/en/training/text_inversion.md @@ -32,11 +32,8 @@ Navigate to the example folder with the training script and install the required cd examples/textual_inversion pip install -r requirements.txt ``` - - -🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It'll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate [Quick tour](https://huggingface.co/docs/accelerate/quicktour) to learn more. - - +> [!TIP] +> 🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It'll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate [Quick tour](https://huggingface.co/docs/accelerate/quicktour) to learn more. Initialize an 🤗 Accelerate environment: @@ -60,11 +57,8 @@ write_basic_config() Lastly, if you want to train a model on your own dataset, take a look at the [Create a dataset for training](create_dataset) guide to learn how to create a dataset that works with the training script. - - -The following sections highlight parts of the training script that are important for understanding how to modify it, but it doesn't cover every aspect of the script in detail. If you're interested in learning more, feel free to read through the [script](https://github.com/huggingface/diffusers/blob/main/examples/textual_inversion/textual_inversion.py) and let us know if you have any questions or concerns. - - +> [!TIP] +> The following sections highlight parts of the training script that are important for understanding how to modify it, but it doesn't cover every aspect of the script in detail. If you're interested in learning more, feel free to read through the [script](https://github.com/huggingface/diffusers/blob/main/examples/textual_inversion/textual_inversion.py) and let us know if you have any questions or concerns. ## Script parameters @@ -160,11 +154,8 @@ Set the environment variable `MODEL_NAME` to a model id on the Hub or a path to - `token_identifier.txt`: the special placeholder token - `type_of_concept.txt`: the type of concept you're training on (either "object" or "style") - - -A full training run takes ~1 hour on a single V100 GPU. - - +> [!WARNING] +> A full training run takes ~1 hour on a single V100 GPU. One more thing before you launch the script. If you're interested in following along with the training process, you can periodically save generated images as training progresses. Add the following parameters to the training command: diff --git a/docs/source/en/training/unconditional_training.md b/docs/source/en/training/unconditional_training.md index d2facc7852..ab3bdd6416 100644 --- a/docs/source/en/training/unconditional_training.md +++ b/docs/source/en/training/unconditional_training.md @@ -31,11 +31,8 @@ cd examples/unconditional_image_generation pip install -r requirements.txt ``` - - -🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It'll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate [Quick tour](https://huggingface.co/docs/accelerate/quicktour) to learn more. - - +> [!TIP] +> 🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It'll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate [Quick tour](https://huggingface.co/docs/accelerate/quicktour) to learn more. Initialize an 🤗 Accelerate environment: @@ -61,11 +58,8 @@ Lastly, if you want to train a model on your own dataset, take a look at the [Cr ## Script parameters - - -The following sections highlight parts of the training script that are important for understanding how to modify it, but it doesn't cover every aspect of the script in detail. If you're interested in learning more, feel free to read through the [script](https://github.com/huggingface/diffusers/blob/main/examples/unconditional_image_generation/train_unconditional.py) and let us know if you have any questions or concerns. - - +> [!TIP] +> The following sections highlight parts of the training script that are important for understanding how to modify it, but it doesn't cover every aspect of the script in detail. If you're interested in learning more, feel free to read through the [script](https://github.com/huggingface/diffusers/blob/main/examples/unconditional_image_generation/train_unconditional.py) and let us know if you have any questions or concerns. The training script provides many parameters to help you customize your training run. All of the parameters and their descriptions are found in the [`parse_args()`](https://github.com/huggingface/diffusers/blob/096f84b05f9514fae9f185cbec0a4d38fbad9919/examples/unconditional_image_generation/train_unconditional.py#L55) function. It provides default values for each parameter, such as the training batch size and learning rate, but you can also set your own values in the training command if you'd like. @@ -163,11 +157,8 @@ Finally, the [training loop](https://github.com/huggingface/diffusers/blob/096f8 Once you've made all your changes or you're okay with the default configuration, you're ready to launch the training script! 🚀 - - -A full training run takes 2 hours on 4xV100 GPUs. - - +> [!WARNING] +> A full training run takes 2 hours on 4xV100 GPUs. diff --git a/docs/source/en/training/wuerstchen.md b/docs/source/en/training/wuerstchen.md index 38a1387dd3..1c362879a6 100644 --- a/docs/source/en/training/wuerstchen.md +++ b/docs/source/en/training/wuerstchen.md @@ -33,11 +33,8 @@ cd examples/wuerstchen/text_to_image pip install -r requirements.txt ``` - - -🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It'll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate [Quick tour](https://huggingface.co/docs/accelerate/quicktour) to learn more. - - +> [!TIP] +> 🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It'll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate [Quick tour](https://huggingface.co/docs/accelerate/quicktour) to learn more. Initialize an 🤗 Accelerate environment: @@ -61,11 +58,8 @@ write_basic_config() Lastly, if you want to train a model on your own dataset, take a look at the [Create a dataset for training](create_dataset) guide to learn how to create a dataset that works with the training script. - - -The following sections highlight parts of the training scripts that are important for understanding how to modify it, but it doesn't cover every aspect of the [script](https://github.com/huggingface/diffusers/blob/main/examples/wuerstchen/text_to_image/train_text_to_image_prior.py) in detail. If you're interested in learning more, feel free to read through the scripts and let us know if you have any questions or concerns. - - +> [!TIP] +> The following sections highlight parts of the training scripts that are important for understanding how to modify it, but it doesn't cover every aspect of the [script](https://github.com/huggingface/diffusers/blob/main/examples/wuerstchen/text_to_image/train_text_to_image_prior.py) in detail. If you're interested in learning more, feel free to read through the scripts and let us know if you have any questions or concerns. ## Script parameters @@ -133,11 +127,8 @@ Once you’ve made all your changes or you’re okay with the default configurat Set the `DATASET_NAME` environment variable to the dataset name from the Hub. This guide uses the [Naruto BLIP captions](https://huggingface.co/datasets/lambdalabs/naruto-blip-captions) dataset, but you can create and train on your own datasets as well (see the [Create a dataset for training](create_dataset) guide). - - -To monitor training progress with Weights & Biases, add the `--report_to=wandb` parameter to the training command. You’ll also need to add the `--validation_prompt` to the training command to keep track of results. This can be really useful for debugging the model and viewing intermediate results. - - +> [!TIP] +> To monitor training progress with Weights & Biases, add the `--report_to=wandb` parameter to the training command. You’ll also need to add the `--validation_prompt` to the training command to keep track of results. This can be really useful for debugging the model and viewing intermediate results. ```bash export DATASET_NAME="lambdalabs/naruto-blip-captions" diff --git a/docs/source/en/tutorials/basic_training.md b/docs/source/en/tutorials/basic_training.md index 9a35b3438f..3aa2ae429b 100644 --- a/docs/source/en/tutorials/basic_training.md +++ b/docs/source/en/tutorials/basic_training.md @@ -18,11 +18,8 @@ Unconditional image generation is a popular application of diffusion models that This tutorial will teach you how to train a [`UNet2DModel`] from scratch on a subset of the [Smithsonian Butterflies](https://huggingface.co/datasets/huggan/smithsonian_butterflies_subset) dataset to generate your own 🦋 butterflies 🦋. - - -💡 This training tutorial is based on the [Training with 🧨 Diffusers](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/training_example.ipynb) notebook. For additional details and context about diffusion models like how they work, check out the notebook! - - +> [!TIP] +> 💡 This training tutorial is based on the [Training with 🧨 Diffusers](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/training_example.ipynb) notebook. For additional details and context about diffusion models like how they work, check out the notebook! Before you begin, make sure you have 🤗 Datasets installed to load and preprocess image datasets, and 🤗 Accelerate, to simplify training on any number of GPUs. The following command will also install [TensorBoard](https://www.tensorflow.org/tensorboard) to visualize training metrics (you can also use [Weights & Biases](https://docs.wandb.ai/) to track your training). @@ -94,11 +91,8 @@ You can easily load the [Smithsonian Butterflies](https://huggingface.co/dataset >>> dataset = load_dataset(config.dataset_name, split="train") ``` - - -💡 You can find additional datasets from the [HugGan Community Event](https://huggingface.co/huggan) or you can use your own dataset by creating a local [`ImageFolder`](https://huggingface.co/docs/datasets/image_dataset#imagefolder). Set `config.dataset_name` to the repository id of the dataset if it is from the HugGan Community Event, or `imagefolder` if you're using your own images. - - +> [!TIP] +> 💡 You can find additional datasets from the [HugGan Community Event](https://huggingface.co/huggan) or you can use your own dataset by creating a local [`ImageFolder`](https://huggingface.co/docs/datasets/image_dataset#imagefolder). Set `config.dataset_name` to the repository id of the dataset if it is from the HugGan Community Event, or `imagefolder` if you're using your own images. 🤗 Datasets uses the [`~datasets.Image`] feature to automatically decode the image data and load it as a [`PIL.Image`](https://pillow.readthedocs.io/en/stable/reference/Image.html) which we can visualize: @@ -274,11 +268,8 @@ Then, you'll need a way to evaluate the model. For evaluation, you can use the [ Now you can wrap all these components together in a training loop with 🤗 Accelerate for easy TensorBoard logging, gradient accumulation, and mixed precision training. To upload the model to the Hub, write a function to get your repository name and information and then push it to the Hub. - - -💡 The training loop below may look intimidating and long, but it'll be worth it later when you launch your training in just one line of code! If you can't wait and want to start generating images, feel free to copy and run the code below. You can always come back and examine the training loop more closely later, like when you're waiting for your model to finish training. 🤗 - - +> [!TIP] +> 💡 The training loop below may look intimidating and long, but it'll be worth it later when you launch your training in just one line of code! If you can't wait and want to start generating images, feel free to copy and run the code below. You can always come back and examine the training loop more closely later, like when you're waiting for your model to finish training. 🤗 ```py >>> from accelerate import Accelerator diff --git a/docs/source/en/using-diffusers/conditional_image_generation.md b/docs/source/en/using-diffusers/conditional_image_generation.md index 7efc0c653e..eb75b6b8a8 100644 --- a/docs/source/en/using-diffusers/conditional_image_generation.md +++ b/docs/source/en/using-diffusers/conditional_image_generation.md @@ -18,11 +18,8 @@ When you think of diffusion models, text-to-image is usually one of the first th From a very high level, a diffusion model takes a prompt and some random initial noise, and iteratively removes the noise to construct an image. The *denoising* process is guided by the prompt, and once the denoising process ends after a predetermined number of time steps, the image representation is decoded into an image. - - -Read the [How does Stable Diffusion work?](https://huggingface.co/blog/stable_diffusion#how-does-stable-diffusion-work) blog post to learn more about how a latent diffusion model works. - - +> [!TIP] +> Read the [How does Stable Diffusion work?](https://huggingface.co/blog/stable_diffusion#how-does-stable-diffusion-work) blog post to learn more about how a latent diffusion model works. You can generate images from a prompt in 🤗 Diffusers in two steps: @@ -176,11 +173,8 @@ image - - -Other models may have different default image sizes depending on the image sizes in the training dataset. For example, SDXL's default image size is 1024x1024 and using lower `height` and `width` values may result in lower quality images. Make sure you check the model's API reference first! - - +> [!WARNING] +> Other models may have different default image sizes depending on the image sizes in the training dataset. For example, SDXL's default image size is 1024x1024 and using lower `height` and `width` values may result in lower quality images. Make sure you check the model's API reference first! ### Guidance scale @@ -272,11 +266,8 @@ There are several ways to exert more control over how an image is generated outs Prompt weighting is a technique for increasing or decreasing the importance of concepts in a prompt to emphasize or minimize certain features in an image. We recommend using the [Compel](https://github.com/damian0815/compel) library to help you generate the weighted prompt embeddings. - - -Learn how to create the prompt embeddings in the [Prompt weighting](weighted_prompts) guide. This example focuses on how to use the prompt embeddings in the pipeline. - - +> [!TIP] +> Learn how to create the prompt embeddings in the [Prompt weighting](weighted_prompts) guide. This example focuses on how to use the prompt embeddings in the pipeline. Once you've created the embeddings, you can pass them to the `prompt_embeds` (and `negative_prompt_embeds` if you're using a negative prompt) parameter in the pipeline. diff --git a/docs/source/en/using-diffusers/controlling_generation.md b/docs/source/en/using-diffusers/controlling_generation.md index 8fd57a7cb8..aed3a8b729 100644 --- a/docs/source/en/using-diffusers/controlling_generation.md +++ b/docs/source/en/using-diffusers/controlling_generation.md @@ -84,23 +84,17 @@ Pix2Pix Zero can be used both to edit synthetic images as well as real images. Next, we generate image captions for the concept that shall be edited and for the new target concept. We can use a model like [Flan-T5](https://huggingface.co/docs/transformers/model_doc/flan-t5) for this purpose. Then, "mean" prompt embeddings for both the source and target concepts are created via the text encoder. Finally, the pix2pix-zero algorithm is used to edit the synthetic image. - To edit a real image, one first generates an image caption using a model like [BLIP](https://huggingface.co/docs/transformers/model_doc/blip). Then one applies DDIM inversion on the prompt and image to generate "inverse" latents. Similar to before, "mean" prompt embeddings for both source and target concepts are created and finally the pix2pix-zero algorithm in combination with the "inverse" latents is used to edit the image. - - -Pix2Pix Zero is the first model that allows "zero-shot" image editing. This means that the model -can edit an image in less than a minute on a consumer GPU as shown [here](../api/pipelines/pix2pix_zero#usage-example). - - +> [!TIP] +> Pix2Pix Zero is the first model that allows "zero-shot" image editing. This means that the model +> can edit an image in less than a minute on a consumer GPU as shown [here](../api/pipelines/pix2pix_zero#usage-example). As mentioned above, Pix2Pix Zero includes optimizing the latents (and not any of the UNet, VAE, or the text encoder) to steer the generation toward a specific concept. This means that the overall pipeline might require more memory than a standard [StableDiffusionPipeline](../api/pipelines/stable_diffusion/text2img). - - -An important distinction between methods like InstructPix2Pix and Pix2Pix Zero is that the former -involves fine-tuning the pre-trained weights while the latter does not. This means that you can -apply Pix2Pix Zero to any of the available Stable Diffusion models. - - +> [!TIP] +> An important distinction between methods like InstructPix2Pix and Pix2Pix Zero is that the former +> involves fine-tuning the pre-trained weights while the latter does not. This means that you can +> apply Pix2Pix Zero to any of the available Stable Diffusion models. ## Attend and Excite diff --git a/docs/source/en/using-diffusers/diffedit.md b/docs/source/en/using-diffusers/diffedit.md index bb1c234dd6..adea210263 100644 --- a/docs/source/en/using-diffusers/diffedit.md +++ b/docs/source/en/using-diffusers/diffedit.md @@ -156,11 +156,8 @@ print(source_prompts) print(target_prompts) ``` - - -Check out the [generation strategy](https://huggingface.co/docs/transformers/main/en/generation_strategies) guide if you're interested in learning more about strategies for generating different quality text. - - +> [!TIP] +> Check out the [generation strategy](https://huggingface.co/docs/transformers/main/en/generation_strategies) guide if you're interested in learning more about strategies for generating different quality text. Load the text encoder model used by the [`StableDiffusionDiffEditPipeline`] to encode the text. You'll use the text encoder to compute the text embeddings: diff --git a/docs/source/en/using-diffusers/img2img.md b/docs/source/en/using-diffusers/img2img.md index 3f42c9396d..ef00bf7f9b 100644 --- a/docs/source/en/using-diffusers/img2img.md +++ b/docs/source/en/using-diffusers/img2img.md @@ -33,11 +33,8 @@ pipeline.enable_model_cpu_offload() pipeline.enable_xformers_memory_efficient_attention() ``` - - -You'll notice throughout the guide, we use [`~DiffusionPipeline.enable_model_cpu_offload`] and [`~DiffusionPipeline.enable_xformers_memory_efficient_attention`], to save memory and increase inference speed. If you're using PyTorch 2.0, then you don't need to call [`~DiffusionPipeline.enable_xformers_memory_efficient_attention`] on your pipeline because it'll already be using PyTorch 2.0's native [scaled-dot product attention](../optimization/fp16#scaled-dot-product-attention). - - +> [!TIP] +> You'll notice throughout the guide, we use [`~DiffusionPipeline.enable_model_cpu_offload`] and [`~DiffusionPipeline.enable_xformers_memory_efficient_attention`], to save memory and increase inference speed. If you're using PyTorch 2.0, then you don't need to call [`~DiffusionPipeline.enable_xformers_memory_efficient_attention`] on your pipeline because it'll already be using PyTorch 2.0's native [scaled-dot product attention](../optimization/fp16#scaled-dot-product-attention). 2. Load an image to pass to the pipeline: @@ -386,11 +383,8 @@ prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k" image = pipeline(prompt, image=init_image, output_type="latent").images[0] ``` - - -It is important to specify `output_type="latent"` in the pipeline to keep all the outputs in latent space to avoid an unnecessary decode-encode step. This only works if the chained pipelines are using the same VAE. - - +> [!TIP] +> It is important to specify `output_type="latent"` in the pipeline to keep all the outputs in latent space to avoid an unnecessary decode-encode step. This only works if the chained pipelines are using the same VAE. Pass the latent output from this pipeline to the next pipeline to generate an image in a [comic book art style](https://huggingface.co/ogkalu/Comic-Diffusion): @@ -449,11 +443,8 @@ prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k" image_1 = pipeline(prompt, image=init_image, output_type="latent").images[0] ``` - - -It is important to specify `output_type="latent"` in the pipeline to keep all the outputs in *latent* space to avoid an unnecessary decode-encode step. This only works if the chained pipelines are using the same VAE. - - +> [!TIP] +> It is important to specify `output_type="latent"` in the pipeline to keep all the outputs in *latent* space to avoid an unnecessary decode-encode step. This only works if the chained pipelines are using the same VAE. Chain it to an upscaler pipeline to increase the image resolution: diff --git a/docs/source/en/using-diffusers/inference_with_tcd_lora.md b/docs/source/en/using-diffusers/inference_with_tcd_lora.md index 88dd4733b5..a4de12b5e7 100644 --- a/docs/source/en/using-diffusers/inference_with_tcd_lora.md +++ b/docs/source/en/using-diffusers/inference_with_tcd_lora.md @@ -335,9 +335,8 @@ grid_image = make_image_grid([canny_image, image], rows=1, cols=2) ``` ![](https://github.com/jabir-zheng/TCD/raw/main/assets/controlnet_canny_tcd.png) - -The inference parameters in this example might not work for all examples, so we recommend you to try different values for `num_inference_steps`, `guidance_scale`, `controlnet_conditioning_scale` and `cross_attention_kwargs` parameters and choose the best one. - +> [!TIP] +> The inference parameters in this example might not work for all examples, so we recommend you to try different values for `num_inference_steps`, `guidance_scale`, `controlnet_conditioning_scale` and `cross_attention_kwargs` parameters and choose the best one. diff --git a/docs/source/en/using-diffusers/inpaint.md b/docs/source/en/using-diffusers/inpaint.md index 695ec04088..28da3a68a5 100644 --- a/docs/source/en/using-diffusers/inpaint.md +++ b/docs/source/en/using-diffusers/inpaint.md @@ -33,11 +33,8 @@ pipeline.enable_model_cpu_offload() pipeline.enable_xformers_memory_efficient_attention() ``` - - -You'll notice throughout the guide, we use [`~DiffusionPipeline.enable_model_cpu_offload`] and [`~DiffusionPipeline.enable_xformers_memory_efficient_attention`], to save memory and increase inference speed. If you're using PyTorch 2.0, it's not necessary to call [`~DiffusionPipeline.enable_xformers_memory_efficient_attention`] on your pipeline because it'll already be using PyTorch 2.0's native [scaled-dot product attention](../optimization/fp16#scaled-dot-product-attention). - - +> [!TIP] +> You'll notice throughout the guide, we use [`~DiffusionPipeline.enable_model_cpu_offload`] and [`~DiffusionPipeline.enable_xformers_memory_efficient_attention`], to save memory and increase inference speed. If you're using PyTorch 2.0, it's not necessary to call [`~DiffusionPipeline.enable_xformers_memory_efficient_attention`] on your pipeline because it'll already be using PyTorch 2.0's native [scaled-dot product attention](../optimization/fp16#scaled-dot-product-attention). 2. Load the base and mask images: @@ -639,11 +636,8 @@ pipeline.enable_xformers_memory_efficient_attention() image = pipeline(prompt=prompt, image=image_inpainting, mask_image=mask_image, output_type="latent").images[0] ``` - - -It is important to specify `output_type="latent"` in the pipeline to keep all the outputs in latent space to avoid an unnecessary decode-encode step. This only works if the chained pipelines are using the same VAE. For example, in the [Text-to-image-to-inpaint](#text-to-image-to-inpaint) section, Kandinsky 2.2 uses a different VAE class than the Stable Diffusion model so it won't work. But if you use Stable Diffusion v1.5 for both pipelines, then you can keep everything in latent space because they both use [`AutoencoderKL`]. - - +> [!TIP] +> It is important to specify `output_type="latent"` in the pipeline to keep all the outputs in latent space to avoid an unnecessary decode-encode step. This only works if the chained pipelines are using the same VAE. For example, in the [Text-to-image-to-inpaint](#text-to-image-to-inpaint) section, Kandinsky 2.2 uses a different VAE class than the Stable Diffusion model so it won't work. But if you use Stable Diffusion v1.5 for both pipelines, then you can keep everything in latent space because they both use [`AutoencoderKL`]. Finally, you can pass this image to an image-to-image pipeline to put the finishing touches on it. It is more efficient to use the [`~AutoPipelineForImage2Image.from_pipe`] method to reuse the existing pipeline components, and avoid unnecessarily loading all the pipeline components into memory again. diff --git a/docs/source/en/using-diffusers/kandinsky.md b/docs/source/en/using-diffusers/kandinsky.md index a482380524..2671c108b3 100644 --- a/docs/source/en/using-diffusers/kandinsky.md +++ b/docs/source/en/using-diffusers/kandinsky.md @@ -31,15 +31,12 @@ Before you begin, make sure you have the following libraries installed: #!pip install -q diffusers transformers accelerate ``` - - -Kandinsky 2.1 and 2.2 usage is very similar! The only difference is Kandinsky 2.2 doesn't accept `prompt` as an input when decoding the latents. Instead, Kandinsky 2.2 only accepts `image_embeds` during decoding. - -
- -Kandinsky 3 has a more concise architecture and it doesn't require a prior model. This means it's usage is identical to other diffusion models like [Stable Diffusion XL](sdxl). - - +> [!WARNING] +> Kandinsky 2.1 and 2.2 usage is very similar! The only difference is Kandinsky 2.2 doesn't accept `prompt` as an input when decoding the latents. Instead, Kandinsky 2.2 only accepts `image_embeds` during decoding. +> +>
+> +> Kandinsky 3 has a more concise architecture and it doesn't require a prior model. This means it's usage is identical to other diffusion models like [Stable Diffusion XL](sdxl). ## Text-to-image @@ -321,20 +318,17 @@ make_image_grid([original_image.resize((512, 512)), image.resize((512, 512))], r ## Inpainting - - -⚠️ The Kandinsky models use ⬜️ **white pixels** to represent the masked area now instead of black pixels. If you are using [`KandinskyInpaintPipeline`] in production, you need to change the mask to use white pixels: - -```py -# For PIL input -import PIL.ImageOps -mask = PIL.ImageOps.invert(mask) - -# For PyTorch and NumPy input -mask = 1 - mask -``` - - +> [!WARNING] +> ⚠️ The Kandinsky models use ⬜️ **white pixels** to represent the masked area now instead of black pixels. If you are using [`KandinskyInpaintPipeline`] in production, you need to change the mask to use white pixels: +> +> ```py +> # For PIL input +> import PIL.ImageOps +> mask = PIL.ImageOps.invert(mask) +> +> # For PyTorch and NumPy input +> mask = 1 - mask +> ``` For inpainting, you'll need the original image, a mask of the area to replace in the original image, and a text prompt of what to inpaint. Load the prior pipeline: @@ -565,11 +559,8 @@ image ## ControlNet - - -⚠️ ControlNet is only supported for Kandinsky 2.2! - - +> [!WARNING] +> ⚠️ ControlNet is only supported for Kandinsky 2.2! ControlNet enables conditioning large pretrained diffusion models with additional inputs such as a depth map or edge detection. For example, you can condition Kandinsky 2.2 with a depth map so the model understands and preserves the structure of the depth image. diff --git a/docs/source/en/using-diffusers/pag.md b/docs/source/en/using-diffusers/pag.md index 46d716bcf8..c11a5dc379 100644 --- a/docs/source/en/using-diffusers/pag.md +++ b/docs/source/en/using-diffusers/pag.md @@ -219,11 +219,8 @@ pipeline = AutoPipelineForText2Image.from_pretrained( pipeline.enable_model_cpu_offload() ``` - - -If you already have a controlnet pipeline and want to enable PAG, you can use the `from_pipe` API: `AutoPipelineForText2Image.from_pipe(pipeline_controlnet, enable_pag=True)` - - +> [!TIP] +> If you already have a controlnet pipeline and want to enable PAG, you can use the `from_pipe` API: `AutoPipelineForText2Image.from_pipe(pipeline_controlnet, enable_pag=True)` You can use the pipeline in the same way you normally use ControlNet pipelines, with the added option to specify a `pag_scale` parameter. Note that PAG works well for unconditional generation. In this example, we will generate an image without a prompt. diff --git a/docs/source/en/using-diffusers/sdxl.md b/docs/source/en/using-diffusers/sdxl.md index 106005c338..79625e0c4a 100644 --- a/docs/source/en/using-diffusers/sdxl.md +++ b/docs/source/en/using-diffusers/sdxl.md @@ -29,15 +29,12 @@ Before you begin, make sure you have the following libraries installed: #!pip install -q diffusers transformers accelerate invisible-watermark>=0.2.0 ``` - - -We recommend installing the [invisible-watermark](https://pypi.org/project/invisible-watermark/) library to help identify images that are generated. If the invisible-watermark library is installed, it is used by default. To disable the watermarker: - -```py -pipeline = StableDiffusionXLPipeline.from_pretrained(..., add_watermarker=False) -``` - - +> [!WARNING] +> We recommend installing the [invisible-watermark](https://pypi.org/project/invisible-watermark/) library to help identify images that are generated. If the invisible-watermark library is installed, it is used by default. To disable the watermarker: +> +> ```py +> pipeline = StableDiffusionXLPipeline.from_pretrained(..., add_watermarker=False) +> ``` ## Load model checkpoints @@ -174,11 +171,8 @@ refiner = DiffusionPipeline.from_pretrained( To use this approach, you need to define the number of timesteps for each model to run through their respective stages. For the base model, this is controlled by the [`denoising_end`](https://huggingface.co/docs/diffusers/main/en/api/pipelines/stable_diffusion/stable_diffusion_xl#diffusers.StableDiffusionXLPipeline.__call__.denoising_end) parameter and for the refiner model, it is controlled by the [`denoising_start`](https://huggingface.co/docs/diffusers/main/en/api/pipelines/stable_diffusion/stable_diffusion_xl#diffusers.StableDiffusionXLImg2ImgPipeline.__call__.denoising_start) parameter. - - -The `denoising_end` and `denoising_start` parameters should be a float between 0 and 1. These parameters are represented as a proportion of discrete timesteps as defined by the scheduler. If you're also using the `strength` parameter, it'll be ignored because the number of denoising steps is determined by the discrete timesteps the model is trained on and the declared fractional cutoff. - - +> [!TIP] +> The `denoising_end` and `denoising_start` parameters should be a float between 0 and 1. These parameters are represented as a proportion of discrete timesteps as defined by the scheduler. If you're also using the `strength` parameter, it'll be ignored because the number of denoising steps is determined by the discrete timesteps the model is trained on and the declared fractional cutoff. Let's set `denoising_end=0.8` so the base model performs the first 80% of denoising the **high-noise** timesteps and set `denoising_start=0.8` so the refiner model performs the last 20% of denoising the **low-noise** timesteps. The base model output should be in **latent** space instead of a PIL image. @@ -285,11 +279,8 @@ refiner = DiffusionPipeline.from_pretrained( ).to("cuda") ``` - - -You can use SDXL refiner with a different base model. For example, you can use the [Hunyuan-DiT](../../api/pipelines/hunyuandit) or [PixArt-Sigma](../../api/pipelines/pixart_sigma) pipelines to generate images with better prompt adherence. Once you have generated an image, you can pass it to the SDXL refiner model to enhance final generation quality. - - +> [!WARNING] +> You can use SDXL refiner with a different base model. For example, you can use the [Hunyuan-DiT](../../api/pipelines/hunyuandit) or [PixArt-Sigma](../../api/pipelines/pixart_sigma) pipelines to generate images with better prompt adherence. Once you have generated an image, you can pass it to the SDXL refiner model to enhance final generation quality. Generate an image from the base model, and set the model output to **latent** space: @@ -322,11 +313,8 @@ For inpainting, load the base and the refiner model in the [`StableDiffusionXLIn SDXL training involves several additional conditioning techniques, which are referred to as *micro-conditioning*. These include original image size, target image size, and cropping parameters. The micro-conditionings can be used at inference time to create high-quality, centered images. - - -You can use both micro-conditioning and negative micro-conditioning parameters thanks to classifier-free guidance. They are available in the [`StableDiffusionXLPipeline`], [`StableDiffusionXLImg2ImgPipeline`], [`StableDiffusionXLInpaintPipeline`], and [`StableDiffusionXLControlNetPipeline`]. - - +> [!TIP] +> You can use both micro-conditioning and negative micro-conditioning parameters thanks to classifier-free guidance. They are available in the [`StableDiffusionXLPipeline`], [`StableDiffusionXLImg2ImgPipeline`], [`StableDiffusionXLInpaintPipeline`], and [`StableDiffusionXLControlNetPipeline`]. ### Size conditioning diff --git a/docs/source/en/using-diffusers/shap-e.md b/docs/source/en/using-diffusers/shap-e.md index 51f0f53b02..8cd62b3ffd 100644 --- a/docs/source/en/using-diffusers/shap-e.md +++ b/docs/source/en/using-diffusers/shap-e.md @@ -151,11 +151,8 @@ images = pipe(prompt, guidance_scale=guidance_scale, num_inference_steps=64, fra Use the [`~utils.export_to_ply`] function to save the mesh output as a `ply` file: - - -You can optionally save the mesh output as an `obj` file with the [`~utils.export_to_obj`] function. The ability to save the mesh output in a variety of formats makes it more flexible for downstream usage! - - +> [!TIP] +> You can optionally save the mesh output as an `obj` file with the [`~utils.export_to_obj`] function. The ability to save the mesh output in a variety of formats makes it more flexible for downstream usage! ```py from diffusers.utils import export_to_ply diff --git a/docs/source/en/using-diffusers/unconditional_image_generation.md b/docs/source/en/using-diffusers/unconditional_image_generation.md index 0208d715d4..0add5bab67 100644 --- a/docs/source/en/using-diffusers/unconditional_image_generation.md +++ b/docs/source/en/using-diffusers/unconditional_image_generation.md @@ -26,11 +26,8 @@ image = generator().images[0] image ``` - - -Want to generate images of something else? Take a look at the training [guide](../training/unconditional_training) to learn how to train a model to generate your own images. - - +> [!TIP] +> Want to generate images of something else? Take a look at the training [guide](../training/unconditional_training) to learn how to train a model to generate your own images. The output image is a [`PIL.Image`](https://pillow.readthedocs.io/en/stable/reference/Image.html?highlight=image#the-image-class) object that can be saved: diff --git a/docs/source/en/using-diffusers/weighted_prompts.md b/docs/source/en/using-diffusers/weighted_prompts.md index 2ebf92d0eb..b45568ac4d 100644 --- a/docs/source/en/using-diffusers/weighted_prompts.md +++ b/docs/source/en/using-diffusers/weighted_prompts.md @@ -217,11 +217,8 @@ Prompt weighting provides a way to emphasize or de-emphasize certain parts of a Prompt weighting works by increasing or decreasing the scale of the text embedding vector that corresponds to its concept in the prompt because you may not necessarily want the model to focus on all concepts equally. The easiest way to prepare the prompt embeddings is to use [Stable Diffusion Long Prompt Weighted Embedding](https://github.com/xhinker/sd_embed) (sd_embed). Once you have the prompt-weighted embeddings, you can pass them to any pipeline that has a [prompt_embeds](https://huggingface.co/docs/diffusers/en/api/pipelines/stable_diffusion/text2img#diffusers.StableDiffusionPipeline.__call__.prompt_embeds) (and optionally [negative_prompt_embeds](https://huggingface.co/docs/diffusers/en/api/pipelines/stable_diffusion/text2img#diffusers.StableDiffusionPipeline.__call__.negative_prompt_embeds)) parameter, such as [`StableDiffusionPipeline`], [`StableDiffusionControlNetPipeline`], and [`StableDiffusionXLPipeline`]. - - -If your favorite pipeline doesn't have a `prompt_embeds` parameter, please open an [issue](https://github.com/huggingface/diffusers/issues/new/choose) so we can add it! - - +> [!TIP] +> If your favorite pipeline doesn't have a `prompt_embeds` parameter, please open an [issue](https://github.com/huggingface/diffusers/issues/new/choose) so we can add it! This guide will show you how to weight your prompts with sd_embed. diff --git a/docs/source/en/using-diffusers/write_own_pipeline.md b/docs/source/en/using-diffusers/write_own_pipeline.md index 15a7e8dc7c..930b0fe21f 100644 --- a/docs/source/en/using-diffusers/write_own_pipeline.md +++ b/docs/source/en/using-diffusers/write_own_pipeline.md @@ -110,11 +110,8 @@ Stable Diffusion is a text-to-image *latent diffusion* model. It is called a lat As you can see, this is already more complex than the DDPM pipeline which only contains a UNet model. The Stable Diffusion model has three separate pretrained models. - - -💡 Read the [How does Stable Diffusion work?](https://huggingface.co/blog/stable_diffusion#how-does-stable-diffusion-work) blog for more details about how the VAE, UNet, and text encoder models work. - - +> [!TIP] +> 💡 Read the [How does Stable Diffusion work?](https://huggingface.co/blog/stable_diffusion#how-does-stable-diffusion-work) blog for more details about how the VAE, UNet, and text encoder models work. Now that you know what you need for the Stable Diffusion pipeline, load all these components with the [`~ModelMixin.from_pretrained`] method. You can find them in the pretrained [`stable-diffusion-v1-5/stable-diffusion-v1-5`](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) checkpoint, and each component is stored in a separate subfolder: @@ -155,11 +152,8 @@ To speed up inference, move the models to a GPU since, unlike the scheduler, the The next step is to tokenize the text to generate embeddings. The text is used to condition the UNet model and steer the diffusion process towards something that resembles the input prompt. - - -💡 The `guidance_scale` parameter determines how much weight should be given to the prompt when generating an image. - - +> [!TIP] +> 💡 The `guidance_scale` parameter determines how much weight should be given to the prompt when generating an image. Feel free to choose any prompt you like if you want to generate something else! @@ -202,15 +196,12 @@ Let's concatenate the conditional and unconditional embeddings into a batch to a Next, generate some initial random noise as a starting point for the diffusion process. This is the latent representation of the image, and it'll be gradually denoised. At this point, the `latent` image is smaller than the final image size but that's okay though because the model will transform it into the final 512x512 image dimensions later. - - -💡 The height and width are divided by 8 because the `vae` model has 3 down-sampling layers. You can check by running the following: - -```py -2 ** (len(vae.config.block_out_channels) - 1) == 8 -``` - - +> [!TIP] +> 💡 The height and width are divided by 8 because the `vae` model has 3 down-sampling layers. You can check by running the following: +> +> ```py +> 2 ** (len(vae.config.block_out_channels) - 1) == 8 +> ``` ```py >>> latents = torch.randn( diff --git a/docs/source/ja/installation.md b/docs/source/ja/installation.md index 97d60528c4..fd6f4eda0f 100644 --- a/docs/source/ja/installation.md +++ b/docs/source/ja/installation.md @@ -108,11 +108,8 @@ pip install -e ".[flax]" Python は通常のライブラリパスに加えて、クローンしたフォルダの中を探すようになります。 例えば、Python パッケージが通常 `~/anaconda3/envs/main/lib/python3.10/site-packages/` にインストールされている場合、Python はクローンした `~/diffusers/` フォルダも同様に参照します。 - - -ライブラリを使い続けたい場合は、`diffusers`フォルダを残しておく必要があります。 - - +> [!WARNING] +> ライブラリを使い続けたい場合は、`diffusers`フォルダを残しておく必要があります。 これで、以下のコマンドで簡単にクローンを最新版の🤗 Diffusersにアップデートできます： diff --git a/docs/source/ja/quicktour.md b/docs/source/ja/quicktour.md index 03b340b352..ce88aaf7b5 100644 --- a/docs/source/ja/quicktour.md +++ b/docs/source/ja/quicktour.md @@ -24,11 +24,8 @@ specific language governing permissions and limitations under the License. この案内では、[`DiffusionPipeline`]を生成に使用する方法を紹介し、モデルとスケジューラを組み合わせて[`DiffusionPipeline`]の内部で起こっていることを再現する方法を説明します。 - - -この案内は🧨 Diffusers [ノートブック](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/diffusers_intro.ipynb)を簡略化したもので、すぐに使い始めることができます。Diffusers 🧨のゴール、設計哲学、コアAPIの詳細についてもっと知りたい方は、ノートブックをご覧ください！ - - +> [!TIP] +> この案内は🧨 Diffusers [ノートブック](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/diffusers_intro.ipynb)を簡略化したもので、すぐに使い始めることができます。Diffusers 🧨のゴール、設計哲学、コアAPIの詳細についてもっと知りたい方は、ノートブックをご覧ください！ 始める前に必要なライブラリーがすべてインストールされていることを確認してください： @@ -56,11 +53,8 @@ specific language governing permissions and limitations under the License. この[`DiffusionPipeline`]はHugging Face Hubに保存されている任意の[チェックポイント](https://huggingface.co/models?library=diffusers&sort=downloads)を使用することができます。 この案内では、[`stable-diffusion-v1-5`](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5)チェックポイントでテキストから画像へ生成します。 - - -[Stable Diffusion]モデルについては、モデルを実行する前にまず[ライセンス](https://huggingface.co/spaces/CompVis/stable-diffusion-license)を注意深くお読みください。🧨 Diffusers は、攻撃的または有害なコンテンツを防ぐために [`safety_checker`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/safety_checker.py) を実装していますが、モデルの改良された画像生成機能により、潜在的に有害なコンテンツが生成される可能性があります。 - - +> [!WARNING] +> [Stable Diffusion]モデルについては、モデルを実行する前にまず[ライセンス](https://huggingface.co/spaces/CompVis/stable-diffusion-license)を注意深くお読みください。🧨 Diffusers は、攻撃的または有害なコンテンツを防ぐために [`safety_checker`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/safety_checker.py) を実装していますが、モデルの改良された画像生成機能により、潜在的に有害なコンテンツが生成される可能性があります。 モデルを[`~DiffusionPipeline.from_pretrained`]メソッドでロードします： @@ -204,11 +198,8 @@ torch.Size([1, 3, 256, 256]) スケジューラは、モデルの出力（この場合は `noisy_residual` ）が与えられたときに、ノイズの多いサンプルからノイズの少ないサンプルへの移行を管理します。 - - -🧨 Diffusersは拡散システムを構築するためのツールボックスです。[`DiffusionPipeline`]は事前に構築された拡散システムを使い始めるのに便利な方法ですが、独自のモデルとスケジューラコンポーネントを個別に選択してカスタム拡散システムを構築することもできます。 - - +> [!TIP] +> 🧨 Diffusersは拡散システムを構築するためのツールボックスです。[`DiffusionPipeline`]は事前に構築された拡散システムを使い始めるのに便利な方法ですが、独自のモデルとスケジューラコンポーネントを個別に選択してカスタム拡散システムを構築することもできます。 この案内では、[`DDPMScheduler`]を[`~diffusers.ConfigMixin.from_config`]メソッドでインスタンス化します： @@ -232,11 +223,8 @@ DDPMScheduler { } ``` - - -💡 スケジューラがどのようにコンフィギュレーションからインスタンス化されるかに注目してください。モデルとは異なり、スケジューラは学習可能な重みを持たず、パラメーターを持ちません！ - - +> [!TIP] +> 💡 スケジューラがどのようにコンフィギュレーションからインスタンス化されるかに注目してください。モデルとは異なり、スケジューラは学習可能な重みを持たず、パラメーターを持ちません！ 最も重要なパラメータは以下の通りです： diff --git a/docs/source/ja/stable_diffusion.md b/docs/source/ja/stable_diffusion.md index 85f2b38a7d..79abfa005d 100644 --- a/docs/source/ja/stable_diffusion.md +++ b/docs/source/ja/stable_diffusion.md @@ -37,11 +37,8 @@ prompt = "portrait photo of a old warrior chief" ## Speed - - -💡 GPUを利用できない場合は、[Colab](https://colab.research.google.com/)のようなGPUプロバイダーから無料で利用できます！ - - +> [!TIP] +> 💡 GPUを利用できない場合は、[Colab](https://colab.research.google.com/)のようなGPUプロバイダーから無料で利用できます！ 画像生成を高速化する最も簡単な方法の1つは、PyTorchモジュールと同じようにGPU上にパイプラインを配置することです： @@ -88,11 +85,8 @@ image 今回、画像生成にかかった時間はわずか11秒で、以前より3倍近く速くなりました！ - - -💡 パイプラインは常に `float16` で実行することを強くお勧めします。 - - +> [!TIP] +> 💡 パイプラインは常に `float16` で実行することを強くお勧めします。 生成ステップ数を減らすという方法もあります。より効率的なスケジューラを選択することで、出力品質を犠牲にすることなくステップ数を減らすことができます。`compatibles`メソッドを呼び出すことで、[`DiffusionPipeline`]の現在のモデルと互換性のあるスケジューラを見つけることができます： diff --git a/docs/source/ja/tutorials/autopipeline.md b/docs/source/ja/tutorials/autopipeline.md index a9a780186a..7dc678da90 100644 --- a/docs/source/ja/tutorials/autopipeline.md +++ b/docs/source/ja/tutorials/autopipeline.md @@ -16,11 +16,8 @@ Diffusersは様々なタスクをこなすことができ、テキストから `AutoPipeline` クラスは、🤗 Diffusers の様々なパイプラインをよりシンプルするために設計されています。この汎用的でタスク重視のパイプラインによってタスクそのものに集中することができます。`AutoPipeline` は、使用するべき正しいパイプラインクラスを自動的に検出するため、特定のパイプラインクラス名を知らなくても、タスクのチェックポイントを簡単にロードできます。 - - -どのタスクがサポートされているかは、[AutoPipeline](../api/pipelines/auto_pipeline) のリファレンスをご覧ください。現在、text-to-image、image-to-image、inpaintingをサポートしています。 - - +> [!TIP] +> どのタスクがサポートされているかは、[AutoPipeline](../api/pipelines/auto_pipeline) のリファレンスをご覧ください。現在、text-to-image、image-to-image、inpaintingをサポートしています。 このチュートリアルでは、`AutoPipeline` を使用して、事前に学習された重みが与えられたときに、特定のタスクを読み込むためのパイプラインクラスを自動的に推測する方法を示します。 diff --git a/docs/source/ko/api/pipelines/stable_diffusion/stable_diffusion_xl.md b/docs/source/ko/api/pipelines/stable_diffusion/stable_diffusion_xl.md index 34a00d63fe..ba85b4a855 100644 --- a/docs/source/ko/api/pipelines/stable_diffusion/stable_diffusion_xl.md +++ b/docs/source/ko/api/pipelines/stable_diffusion/stable_diffusion_xl.md @@ -207,11 +207,8 @@ image = refiner( 동일한 40 단계에서 base 모델을 실행한다면, 이미지의 디테일(예: 사자의 눈과 코)이 떨어졌을 것입니다: - - -앙상블 방식은 사용 가능한 모든 스케줄러에서 잘 작동합니다! - - +> [!TIP] +> 앙상블 방식은 사용 가능한 모든 스케줄러에서 잘 작동합니다! #### 2.) 노이즈가 완전히 제거된 기본 이미지에서 이미지 출력을 정제하기 @@ -248,11 +245,8 @@ image = refiner(prompt=prompt, image=image[None, :]).images[0] |---|---| | ![](https://huggingface.co/datasets/diffusers/docs-images/resolve/main/sd_xl/init_image.png) | ![](https://huggingface.co/datasets/diffusers/docs-images/resolve/main/sd_xl/refined_image.png) | - - -refiner는 또한 인페인팅 설정에 잘 사용될 수 있습니다. 아래에 보여지듯이 [`StableDiffusionXLInpaintPipeline`] 클래스를 사용해서 만들어보세요. - - +> [!TIP] +> refiner는 또한 인페인팅 설정에 잘 사용될 수 있습니다. 아래에 보여지듯이 [`StableDiffusionXLInpaintPipeline`] 클래스를 사용해서 만들어보세요. Denoiser 앙상블 설정에서 인페인팅에 refiner를 사용하려면 다음을 수행하면 됩니다: diff --git a/docs/source/ko/conceptual/evaluation.md b/docs/source/ko/conceptual/evaluation.md index 2d296420bc..731b511485 100644 --- a/docs/source/ko/conceptual/evaluation.md +++ b/docs/source/ko/conceptual/evaluation.md @@ -95,11 +95,8 @@ images = sd_pipeline(sample_prompts, num_images_per_prompt=1, generator=generato 다양한 모델을 사용하여 모든 프롬프트에서 생성된 여러 이미지들이 생성되면 (평가 과정에서) 이러한 결과물들은 사람 평가자들에게 점수를 매기기 위해 제시됩니다. DrawBench와 PartiPrompts 벤치마크에 대한 자세한 내용은 각각의 논문을 참조하십시오. - - -모델이 훈련 중일 때 추론 샘플을 살펴보는 것은 훈련 진행 상황을 측정하는 데 유용합니다. [훈련 스크립트](https://github.com/huggingface/diffusers/tree/main/examples/)에서는 TensorBoard와 Weights & Biases에 대한 추가 지원과 함께 이 유틸리티를 지원합니다. - - +> [!TIP] +> 모델이 훈련 중일 때 추론 샘플을 살펴보는 것은 훈련 진행 상황을 측정하는 데 유용합니다. [훈련 스크립트](https://github.com/huggingface/diffusers/tree/main/examples/)에서는 TensorBoard와 Weights & Biases에 대한 추가 지원과 함께 이 유틸리티를 지원합니다. ## 정량적 평가[[quantitative-evaluation]] @@ -193,11 +190,8 @@ print(f"CLIP Score with v-1-5: {sd_clip_score_1_5}") [v1-5](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) 체크포인트가 이전 버전보다 더 나은 성능을 보이는 것 같습니다. 그러나 CLIP 점수를 계산하기 위해 사용한 프롬프트의 수가 상당히 적습니다. 보다 실용적인 평가를 위해서는 이 수를 훨씬 높게 설정하고, 프롬프트를 다양하게 사용해야 합니다. - - -이 점수에는 몇 가지 제한 사항이 있습니다. 훈련 데이터셋의 캡션은 웹에서 크롤링되어 이미지와 관련된 `alt` 및 유사한 태그에서 추출되었습니다. 이들은 인간이 이미지를 설명하는 데 사용할 수 있는 것과 일치하지 않을 수 있습니다. 따라서 여기서는 몇 가지 프롬프트를 "엔지니어링"해야 했습니다. - - +> [!WARNING] +> 이 점수에는 몇 가지 제한 사항이 있습니다. 훈련 데이터셋의 캡션은 웹에서 크롤링되어 이미지와 관련된 `alt` 및 유사한 태그에서 추출되었습니다. 이들은 인간이 이미지를 설명하는 데 사용할 수 있는 것과 일치하지 않을 수 있습니다. 따라서 여기서는 몇 가지 프롬프트를 "엔지니어링"해야 했습니다. ### 이미지 조건화된 텍스트-이미지 생성[[image-conditioned-text-to-image-generation]] @@ -405,11 +399,8 @@ CLIP 점수와 마찬가지로, CLIP 방향 유사성이 높을수록 좋습니 [`StableDiffusionPix2PixZeroPipeline`](https://huggingface.co/docs/diffusers/main/en/api/pipelines/pix2pix_zero#diffusers.StableDiffusionPix2PixZeroPipeline)와 같은 유사한 파이프라인에도 이러한 메트릭을 사용할 수 있습니다. - - -CLIP 점수와 CLIP 방향 유사성 모두 CLIP 모델에 의존하기 때문에 평가가 편향될 수 있습니다 - - +> [!TIP] +> CLIP 점수와 CLIP 방향 유사성 모두 CLIP 모델에 의존하기 때문에 평가가 편향될 수 있습니다 ***IS, FID (나중에 설명할 예정), 또는 KID와 같은 메트릭을 확장하는 것은 어려울 수 있습니다***. 평가 중인 모델이 대규모 이미지 캡셔닝 데이터셋 (예: [LAION-5B 데이터셋](https://laion.ai/blog/laion-5b/))에서 사전 훈련되었을 때 이는 문제가 될 수 있습니다. 왜냐하면 이러한 메트릭의 기반에는 중간 이미지 특징을 추출하기 위해 ImageNet-1k 데이터셋에서 사전 훈련된 InceptionNet이 사용되기 때문입니다. Stable Diffusion의 사전 훈련 데이터셋은 InceptionNet의 사전 훈련 데이터셋과 겹치는 부분이 제한적일 수 있으므로 따라서 여기에는 좋은 후보가 아닙니다. @@ -532,19 +523,16 @@ FID는 낮을수록 좋습니다. 여러 가지 요소가 FID에 영향을 줄 마지막 두 가지 요소에 대해서는, 다른 시드와 추론 단계에서 평가를 실행하고 평균 결과를 보고하는 것은 좋은 실천 방법입니다 - - -FID 결과는 많은 요소에 의존하기 때문에 취약할 수 있습니다: - -* 계산 중 사용되는 특정 Inception 모델. -* 계산의 구현 정확도. -* 이미지 형식 (PNG 또는 JPG에서 시작하는 경우가 다릅니다). - -이러한 사항을 염두에 두면, FID는 유사한 실행을 비교할 때 가장 유용하지만, 저자가 FID 측정 코드를 주의 깊게 공개하지 않는 한 논문 결과를 재현하기는 어렵습니다. - -이러한 사항은 KID 및 IS와 같은 다른 관련 메트릭에도 적용됩니다. - - +> [!WARNING] +> FID 결과는 많은 요소에 의존하기 때문에 취약할 수 있습니다: +> +> * 계산 중 사용되는 특정 Inception 모델. +> * 계산의 구현 정확도. +> * 이미지 형식 (PNG 또는 JPG에서 시작하는 경우가 다릅니다). +> +> 이러한 사항을 염두에 두면, FID는 유사한 실행을 비교할 때 가장 유용하지만, 저자가 FID 측정 코드를 주의 깊게 공개하지 않는 한 논문 결과를 재현하기는 어렵습니다. +> +> 이러한 사항은 KID 및 IS와 같은 다른 관련 메트릭에도 적용됩니다. 마지막 단계로, `fake_images`를 시각적으로 검사해 봅시다. diff --git a/docs/source/ko/installation.md b/docs/source/ko/installation.md index c03b464290..198ca4b7c7 100644 --- a/docs/source/ko/installation.md +++ b/docs/source/ko/installation.md @@ -107,11 +107,8 @@ pip install -e ".[flax]" Python은 이제 일반 라이브러리 경로에 더하여 복제한 폴더 내부를 살펴봅니다. 예를들어 Python 패키지가 `~/anaconda3/envs/main/lib/python3.10/site-packages/`에 설치되어 있는 경우 Python은 복제한 폴더인 `~/diffusers/`도 검색합니다. - - -라이브러리를 계속 사용하려면 `diffusers` 폴더를 유지해야 합니다. - - +> [!WARNING] +> 라이브러리를 계속 사용하려면 `diffusers` 폴더를 유지해야 합니다. 이제 다음 명령어를 사용하여 최신 버전의 🤗 Diffusers로 쉽게 업데이트할 수 있습니다: diff --git a/docs/source/ko/optimization/coreml.md b/docs/source/ko/optimization/coreml.md index 60f19fd2c3..73ca851177 100644 --- a/docs/source/ko/optimization/coreml.md +++ b/docs/source/ko/optimization/coreml.md @@ -16,11 +16,8 @@ specific language governing permissions and limitations under the License. Core ML 모델은 Apple 기기에서 사용할 수 있는 모든 컴퓨팅 엔진들, 즉 CPU, GPU, Apple Neural Engine(또는 Apple Silicon Mac 및 최신 iPhone/iPad에서 사용할 수 있는 텐서 최적화 가속기인 ANE)을 활용할 수 있습니다. 모델과 실행 중인 기기에 따라 Core ML은 컴퓨팅 엔진도 혼합하여 사용할 수 있으므로, 예를 들어 모델의 일부가 CPU에서 실행되는 반면 다른 부분은 GPU에서 실행될 수 있습니다. - - -PyTorch에 내장된 `mps` 가속기를 사용하여 Apple Silicon Macs에서 `diffusers` Python 코드베이스를 실행할 수도 있습니다. 이 방법은 [mps 가이드]에 자세히 설명되어 있지만 네이티브 앱과 호환되지 않습니다. - - +> [!TIP] +> PyTorch에 내장된 `mps` 가속기를 사용하여 Apple Silicon Macs에서 `diffusers` Python 코드베이스를 실행할 수도 있습니다. 이 방법은 [mps 가이드]에 자세히 설명되어 있지만 네이티브 앱과 호환되지 않습니다. ## Stable Diffusion Core ML 체크포인트 diff --git a/docs/source/ko/optimization/fp16.md b/docs/source/ko/optimization/fp16.md index db0370875e..56f1330c40 100644 --- a/docs/source/ko/optimization/fp16.md +++ b/docs/source/ko/optimization/fp16.md @@ -74,18 +74,16 @@ prompt = "a photo of an astronaut riding a horse on mars" image = pipe(prompt).images[0] ``` - - 어떤 파이프라인에서도 [`torch.autocast`](https://pytorch.org/docs/stable/amp.html#torch.autocast) 를 사용하는 것은 검은색 이미지를 생성할 수 있고, 순수한 float16 정밀도를 사용하는 것보다 항상 느리기 때문에 사용하지 않는 것이 좋습니다. - +> [!WARNING] +> 어떤 파이프라인에서도 [`torch.autocast`](https://pytorch.org/docs/stable/amp.html#torch.autocast) 를 사용하는 것은 검은색 이미지를 생성할 수 있고, 순수한 float16 정밀도를 사용하는 것보다 항상 느리기 때문에 사용하지 않는 것이 좋습니다. ## 추가 메모리 절약을 위한 슬라이스 어텐션 추가 메모리 절약을 위해, 한 번에 모두 계산하는 대신 단계적으로 계산을 수행하는 슬라이스 버전의 어텐션(attention)을 사용할 수 있습니다. - - Attention slicing은 모델이 하나 이상의 어텐션 헤드를 사용하는 한, 배치 크기가 1인 경우에도 유용합니다. - 하나 이상의 어텐션 헤드가 있는 경우 *QK^T* 어텐션 매트릭스는 상당한 양의 메모리를 절약할 수 있는 각 헤드에 대해 순차적으로 계산될 수 있습니다. - +> [!TIP] +> Attention slicing은 모델이 하나 이상의 어텐션 헤드를 사용하는 한, 배치 크기가 1인 경우에도 유용합니다. +> 하나 이상의 어텐션 헤드가 있는 경우 *QK^T* 어텐션 매트릭스는 상당한 양의 메모리를 절약할 수 있는 각 헤드에 대해 순차적으로 계산될 수 있습니다. 각 헤드에 대해 순차적으로 어텐션 계산을 수행하려면, 다음과 같이 추론 전에 파이프라인에서 [`~StableDiffusionPipeline.enable_attention_slicing`]를 호출하면 됩니다: @@ -161,9 +159,8 @@ image = pipe(prompt).images[0] 참고로 이 방법은 전체 모델이 아닌 서브모듈 수준에서 작동합니다. 이는 메모리 소비를 최소화하는 가장 좋은 방법이지만 프로세스의 반복적 특성으로 인해 추론 속도가 훨씬 느립니다. 파이프라인의 UNet 구성 요소는 여러 번 실행됩니다('num_inference_steps' 만큼). 매번 UNet의 서로 다른 서브모듈이 순차적으로 온로드된 다음 필요에 따라 오프로드되므로 메모리 이동 횟수가 많습니다. - -또 다른 최적화 방법인 모델 오프로딩을 사용하는 것을 고려하십시오. 이는 훨씬 빠르지만 메모리 절약이 크지는 않습니다. - +> [!TIP] +> 또 다른 최적화 방법인 모델 오프로딩을 사용하는 것을 고려하십시오. 이는 훨씬 빠르지만 메모리 절약이 크지는 않습니다. 또한 ttention slicing과 연결해서 최소 메모리(< 2GB)로도 동작할 수 있습니다. @@ -231,9 +228,8 @@ pipe.enable_attention_slicing(1) image = pipe(prompt).images[0] ``` - -이 기능을 사용하려면 'accelerate' 버전 0.17.0 이상이 필요합니다. - +> [!TIP] +> 이 기능을 사용하려면 'accelerate' 버전 0.17.0 이상이 필요합니다. ## Channels Last 메모리 형식 사용하기 diff --git a/docs/source/ko/optimization/mps.md b/docs/source/ko/optimization/mps.md index 4daeaf5dba..004374c4af 100644 --- a/docs/source/ko/optimization/mps.md +++ b/docs/source/ko/optimization/mps.md @@ -27,11 +27,8 @@ Diffusers는 Stable Diffusion 추론을 위해 PyTorch `mps`를 사용해 Apple 아래 코도는 익숙한 `to()` 인터페이스를 사용하여 `mps` 백엔드로 Stable Diffusion 파이프라인을 M1 또는 M2 장치로 이동하는 방법을 보여줍니다. - - -**PyTorch 1.13을 사용 중일 때 ** 추가 일회성 전달을 사용하여 파이프라인을 "프라이밍"하는 것을 추천합니다. 이것은 발견한 이상한 문제에 대한 임시 해결 방법입니다. 첫 번째 추론 전달은 후속 전달와 약간 다른 결과를 생성합니다. 이 전달은 한 번만 수행하면 되며 추론 단계를 한 번만 사용하고 결과를 폐기해도 됩니다. - - +> [!WARNING] +> **PyTorch 1.13을 사용 중일 때 ** 추가 일회성 전달을 사용하여 파이프라인을 "프라이밍"하는 것을 추천합니다. 이것은 발견한 이상한 문제에 대한 임시 해결 방법입니다. 첫 번째 추론 전달은 후속 전달와 약간 다른 결과를 생성합니다. 이 전달은 한 번만 수행하면 되며 추론 단계를 한 번만 사용하고 결과를 폐기해도 됩니다. 이전 팁에서 설명한 것들을 포함한 여러 문제를 해결하므로 PyTorch 2 이상을 사용하는 것이 좋습니다. diff --git a/docs/source/ko/optimization/xformers.md b/docs/source/ko/optimization/xformers.md index 3e4d107c0a..96fab34acf 100644 --- a/docs/source/ko/optimization/xformers.md +++ b/docs/source/ko/optimization/xformers.md @@ -21,16 +21,10 @@ specific language governing permissions and limitations under the License. pip install xformers ``` - - -xFormers PIP 패키지에는 최신 버전의 PyTorch(xFormers 0.0.16에 1.13.1)가 필요합니다. 이전 버전의 PyTorch를 사용해야 하는 경우 [프로젝트 지침](https://github.com/facebookresearch/xformers#installing-xformers)의 소스를 사용해 xFormers를 설치하는 것이 좋습니다. - - +> [!TIP] +> xFormers PIP 패키지에는 최신 버전의 PyTorch(xFormers 0.0.16에 1.13.1)가 필요합니다. 이전 버전의 PyTorch를 사용해야 하는 경우 [프로젝트 지침](https://github.com/facebookresearch/xformers#installing-xformers)의 소스를 사용해 xFormers를 설치하는 것이 좋습니다. xFormers를 설치하면, [여기](fp16#memory-efficient-attention)서 설명한 것처럼 'enable_xformers_memory_efficient_attention()'을 사용하여 추론 속도를 높이고 메모리 소비를 줄일 수 있습니다. - - -[이 이슈](https://github.com/huggingface/diffusers/issues/2234#issuecomment-1416931212)에 따르면 xFormers `v0.0.16`에서 GPU를 사용한 학습(파인 튜닝 또는 Dreambooth)을 할 수 없습니다. 해당 문제가 발견되면. 해당 코멘트를 참고해 development 버전을 설치하세요. - - +> [!WARNING] +> [이 이슈](https://github.com/huggingface/diffusers/issues/2234#issuecomment-1416931212)에 따르면 xFormers `v0.0.16`에서 GPU를 사용한 학습(파인 튜닝 또는 Dreambooth)을 할 수 없습니다. 해당 문제가 발견되면. 해당 코멘트를 참고해 development 버전을 설치하세요. diff --git a/docs/source/ko/quicktour.md b/docs/source/ko/quicktour.md index 58ebb8960f..0a3cd0f7c4 100644 --- a/docs/source/ko/quicktour.md +++ b/docs/source/ko/quicktour.md @@ -23,11 +23,8 @@ Diffusion 모델은 이미지나 오디오와 같은 관심 샘플들을 생성 훑어보기에서는 추론을 위해 [`DiffusionPipeline`]을 사용하는 방법을 보여준 다음, 모델과 스케줄러를 결합하여 [`DiffusionPipeline`] 내부에서 일어나는 일을 복제하는 방법을 안내합니다. - - -훑어보기는 간결한 버전의 🧨 Diffusers 소개로서 [노트북](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/diffusers_intro.ipynb) 빠르게 시작할 수 있도록 도와드립니다. 디퓨저의 목표, 디자인 철학, 핵심 API에 대한 추가 세부 정보를 자세히 알아보려면 노트북을 확인하세요! - - +> [!TIP] +> 훑어보기는 간결한 버전의 🧨 Diffusers 소개로서 [노트북](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/diffusers_intro.ipynb) 빠르게 시작할 수 있도록 도와드립니다. 디퓨저의 목표, 디자인 철학, 핵심 API에 대한 추가 세부 정보를 자세히 알아보려면 노트북을 확인하세요! 시작하기 전에 필요한 라이브러리가 모두 설치되어 있는지 확인하세요: @@ -55,11 +52,8 @@ Diffusion 모델은 이미지나 오디오와 같은 관심 샘플들을 생성 허깅페이스 허브에 저장된 모든 [checkpoint](https://huggingface.co/models?library=diffusers&sort=downloads)에 대해 [`DiffusionPipeline`]을 사용할 수 있습니다. 이 훑어보기에서는 text-to-image 생성을 위한 [`stable-diffusion-v1-5`](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) 체크포인트를 로드합니다. - - -[Stable Diffusion](https://huggingface.co/CompVis/stable-diffusion) 모델의 경우, 모델을 실행하기 전에 [라이선스](https://huggingface.co/spaces/CompVis/stable-diffusion-license)를 먼저 주의 깊게 읽어주세요. 🧨 Diffusers는 불쾌하거나 유해한 콘텐츠를 방지하기 위해 [`safety_checker`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/safety_checker.py)를 구현하고 있지만, 모델의 향상된 이미지 생성 기능으로 인해 여전히 잠재적으로 유해한 콘텐츠가 생성될 수 있습니다. - - +> [!WARNING] +> [Stable Diffusion](https://huggingface.co/CompVis/stable-diffusion) 모델의 경우, 모델을 실행하기 전에 [라이선스](https://huggingface.co/spaces/CompVis/stable-diffusion-license)를 먼저 주의 깊게 읽어주세요. 🧨 Diffusers는 불쾌하거나 유해한 콘텐츠를 방지하기 위해 [`safety_checker`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/safety_checker.py)를 구현하고 있지만, 모델의 향상된 이미지 생성 기능으로 인해 여전히 잠재적으로 유해한 콘텐츠가 생성될 수 있습니다. [`~DiffusionPipeline.from_pretrained`] 방법으로 모델 로드하기: @@ -203,11 +197,8 @@ torch.Size([1, 3, 256, 256]) 스케줄러는 모델 출력이 주어졌을 때 노이즈가 많은 샘플에서 노이즈가 적은 샘플로 전환하는 것을 관리합니다 - 이 경우 'noisy_residual'. - - -🧨 Diffusers는 Diffusion 시스템을 구축하기 위한 툴박스입니다. [`DiffusionPipeline`]을 사용하면 미리 만들어진 Diffusion 시스템을 편리하게 시작할 수 있지만, 모델과 스케줄러 구성 요소를 개별적으로 선택하여 사용자 지정 Diffusion 시스템을 구축할 수도 있습니다. - - +> [!TIP] +> 🧨 Diffusers는 Diffusion 시스템을 구축하기 위한 툴박스입니다. [`DiffusionPipeline`]을 사용하면 미리 만들어진 Diffusion 시스템을 편리하게 시작할 수 있지만, 모델과 스케줄러 구성 요소를 개별적으로 선택하여 사용자 지정 Diffusion 시스템을 구축할 수도 있습니다. 훑어보기의 경우, [`~diffusers.ConfigMixin.from_config`] 메서드를 사용하여 [`DDPMScheduler`]를 인스턴스화합니다: @@ -231,11 +222,8 @@ DDPMScheduler { } ``` - - -💡 스케줄러가 구성에서 어떻게 인스턴스화되는지 주목하세요. 모델과 달리 스케줄러에는 학습 가능한 가중치가 없으며 매개변수도 없습니다! - - +> [!TIP] +> 💡 스케줄러가 구성에서 어떻게 인스턴스화되는지 주목하세요. 모델과 달리 스케줄러에는 학습 가능한 가중치가 없으며 매개변수도 없습니다! 가장 중요한 매개변수는 다음과 같습니다: diff --git a/docs/source/ko/stable_diffusion.md b/docs/source/ko/stable_diffusion.md index 794bdf9c66..0f61e16d2a 100644 --- a/docs/source/ko/stable_diffusion.md +++ b/docs/source/ko/stable_diffusion.md @@ -37,11 +37,8 @@ prompt = "portrait photo of a old warrior chief" ## 속도 - - -💡 GPU에 액세스할 수 없는 경우 다음과 같은 GPU 제공업체에서 무료로 사용할 수 있습니다!. [Colab](https://colab.research.google.com/) - - +> [!TIP] +> 💡 GPU에 액세스할 수 없는 경우 다음과 같은 GPU 제공업체에서 무료로 사용할 수 있습니다!. [Colab](https://colab.research.google.com/) 추론 속도를 높이는 가장 간단한 방법 중 하나는 Pytorch 모듈을 사용할 때와 같은 방식으로 GPU에 파이프라인을 배치하는 것입니다: @@ -89,11 +86,8 @@ image 이번에는 이미지를 생성하는 데 약 11초밖에 걸리지 않아 이전보다 3배 가까이 빨라졌습니다! - - -💡 파이프라인은 항상 `float16`에서 실행할 것을 강력히 권장하며, 지금까지 출력 품질이 저하되는 경우는 거의 없었습니다. - - +> [!TIP] +> 💡 파이프라인은 항상 `float16`에서 실행할 것을 강력히 권장하며, 지금까지 출력 품질이 저하되는 경우는 거의 없었습니다. 또 다른 옵션은 추론 단계의 수를 줄이는 것입니다. 보다 효율적인 스케줄러를 선택하면 출력 품질 저하 없이 단계 수를 줄이는 데 도움이 될 수 있습니다. 현재 모델과 호환되는 스케줄러는 `compatibles` 메서드를 호출하여 [`DiffusionPipeline`]에서 찾을 수 있습니다: diff --git a/docs/source/ko/training/controlnet.md b/docs/source/ko/training/controlnet.md index 434ca959bd..e868b57c55 100644 --- a/docs/source/ko/training/controlnet.md +++ b/docs/source/ko/training/controlnet.md @@ -20,11 +20,8 @@ specific language governing permissions and limitations under the License. 아래의 스크립트를 실행하기 전에, 라이브러리의 학습 의존성을 설치해야 합니다. - - -가장 최신 버전의 예시 스크립트를 성공적으로 실행하기 위해서는, 소스에서 설치하고 최신 버전의 설치를 유지하는 것을 강력하게 추천합니다. 우리는 예시 스크립트들을 자주 업데이트하고 예시에 맞춘 특정한 요구사항을 설치합니다. - - +> [!WARNING] +> 가장 최신 버전의 예시 스크립트를 성공적으로 실행하기 위해서는, 소스에서 설치하고 최신 버전의 설치를 유지하는 것을 강력하게 추천합니다. 우리는 예시 스크립트들을 자주 업데이트하고 예시에 맞춘 특정한 요구사항을 설치합니다. 위 사항을 만족시키기 위해서, 새로운 가상환경에서 다음 일련의 스텝을 실행하세요: diff --git a/docs/source/ko/training/create_dataset.md b/docs/source/ko/training/create_dataset.md index a869cd09f0..c459a9d6a1 100644 --- a/docs/source/ko/training/create_dataset.md +++ b/docs/source/ko/training/create_dataset.md @@ -11,11 +11,8 @@ - 이미지 폴더를 `--train_data_dir` 인수에 제공합니다. - 데이터셋을 Hub에 업로드하고 데이터셋 리포지토리 id를 `--dataset_name` 인수에 전달합니다. - - -💡 학습에 사용할 이미지 데이터셋을 만드는 방법에 대한 자세한 내용은 [이미지 데이터셋 만들기](https://huggingface.co/docs/datasets/image_dataset) 가이드를 참고하세요. - - +> [!TIP] +> 💡 학습에 사용할 이미지 데이터셋을 만드는 방법에 대한 자세한 내용은 [이미지 데이터셋 만들기](https://huggingface.co/docs/datasets/image_dataset) 가이드를 참고하세요. ## 폴더 형태로 데이터셋 구축하기 @@ -40,11 +37,8 @@ accelerate launch train_unconditional.py \ ## Hub에 데이터 올리기 - - -💡 데이터셋을 만들고 Hub에 업로드하는 것에 대한 자세한 내용은 [🤗 Datasets을 사용한 이미지 검색](https://huggingface.co/blog/image-search-datasets) 게시물을 참고하세요. - - +> [!TIP] +> 💡 데이터셋을 만들고 Hub에 업로드하는 것에 대한 자세한 내용은 [🤗 Datasets을 사용한 이미지 검색](https://huggingface.co/blog/image-search-datasets) 게시물을 참고하세요. PIL 인코딩된 이미지가 포함된 `이미지` 열을 생성하는 [이미지 폴더](https://huggingface.co/docs/datasets/image_load#imagefolder) 기능을 사용하여 데이터셋 생성을 시작합니다. diff --git a/docs/source/ko/training/distributed_inference.md b/docs/source/ko/training/distributed_inference.md index c4d6400d97..e63764f5eb 100644 --- a/docs/source/ko/training/distributed_inference.md +++ b/docs/source/ko/training/distributed_inference.md @@ -32,9 +32,8 @@ Use the `--num_processes` argument to specify the number of GPUs to use, and cal accelerate launch run_distributed.py --num_processes=2 ``` -자세한 내용은 [🤗 Accelerate를 사용한 분산 추론](https://huggingface.co/docs/accelerate/en/usage_guides/distributed_inference#distributed-inference-with-accelerate) 가이드를 참조하세요. - - +> [!TIP] +> 자세한 내용은 [🤗 Accelerate를 사용한 분산 추론](https://huggingface.co/docs/accelerate/en/usage_guides/distributed_inference#distributed-inference-with-accelerate) 가이드를 참조하세요. ## Pytoerch 분산 diff --git a/docs/source/ko/training/dreambooth.md b/docs/source/ko/training/dreambooth.md index 8e62f8edab..3e5a17d5f6 100644 --- a/docs/source/ko/training/dreambooth.md +++ b/docs/source/ko/training/dreambooth.md @@ -51,11 +51,8 @@ write_basic_config() ## 파인튜닝 - - -DreamBooth 파인튜닝은 하이퍼파라미터에 매우 민감하고 과적합되기 쉽습니다. 적절한 하이퍼파라미터를 선택하는 데 도움이 되도록 다양한 권장 설정이 포함된 [심층 분석](https://huggingface.co/blog/dreambooth)을 살펴보는 것이 좋습니다. - - +> [!WARNING] +> DreamBooth 파인튜닝은 하이퍼파라미터에 매우 민감하고 과적합되기 쉽습니다. 적절한 하이퍼파라미터를 선택하는 데 도움이 되도록 다양한 권장 설정이 포함된 [심층 분석](https://huggingface.co/blog/dreambooth)을 살펴보는 것이 좋습니다. @@ -176,11 +173,8 @@ python train_dreambooth_flax.py \ 해당 스크립트를 사용하면 `unet`과 함께 `text_encoder`를 파인튜닝할 수 있습니다. 실험에서(자세한 내용은 [🧨 Diffusers를 사용해 DreamBooth로 Stable Diffusion 학습하기](https://huggingface.co/blog/dreambooth) 게시물을 확인하세요), 특히 얼굴 이미지를 생성할 때 훨씬 더 나은 결과를 얻을 수 있습니다. - - -텍스트 인코더를 학습시키려면 추가 메모리가 필요해 16GB GPU로는 동작하지 않습니다. 이 옵션을 사용하려면 최소 24GB VRAM이 필요합니다. - - +> [!WARNING] +> 텍스트 인코더를 학습시키려면 추가 메모리가 필요해 16GB GPU로는 동작하지 않습니다. 이 옵션을 사용하려면 최소 24GB VRAM이 필요합니다. `--train_text_encoder` 인수를 학습 스크립트에 전달하여 `text_encoder` 및 `unet`을 파인튜닝할 수 있습니다: diff --git a/docs/source/ko/training/lora.md b/docs/source/ko/training/lora.md index 5bcef27143..515e3fd65e 100644 --- a/docs/source/ko/training/lora.md +++ b/docs/source/ko/training/lora.md @@ -14,11 +14,8 @@ specific language governing permissions and limitations under the License. [[open-in-colab]] - - -현재 LoRA는 [`UNet2DConditionalModel`]의 어텐션 레이어에서만 지원됩니다. - - +> [!WARNING] +> 현재 LoRA는 [`UNet2DConditionalModel`]의 어텐션 레이어에서만 지원됩니다. [LoRA(Low-Rank Adaptation of Large Language Models)](https://huggingface.co/papers/2106.09685)는 메모리를 적게 사용하면서 대규모 모델의 학습을 가속화하는 학습 방법입니다. 이는 rank-decomposition weight 행렬 쌍(**업데이트 행렬**이라고 함)을 추가하고 새로 추가된 가중치**만** 학습합니다. 여기에는 몇 가지 장점이 있습니다. @@ -28,11 +25,8 @@ specific language governing permissions and limitations under the License. - 메모리 효율성이 향상되어 Tesla T4, RTX 3080 또는 RTX 2080 Ti와 같은 소비자용 GPU에서 파인튜닝을 실행할 수 있습니다! T4와 같은 GPU는 무료이며 Kaggle 또는 Google Colab 노트북에서 쉽게 액세스할 수 있습니다. - - -💡 LoRA는 어텐션 레이어에만 한정되지는 않습니다. 저자는 언어 모델의 어텐션 레이어를 수정하는 것이 매우 효율적으로 죻은 성능을 얻기에 충분하다는 것을 발견했습니다. 이것이 LoRA 가중치를 모델의 어텐션 레이어에 추가하는 것이 일반적인 이유입니다. LoRA 작동 방식에 대한 자세한 내용은 [Using LoRA for effective Stable Diffusion fine-tuning](https://huggingface.co/blog/lora) 블로그를 확인하세요! - - +> [!TIP] +> 💡 LoRA는 어텐션 레이어에만 한정되지는 않습니다. 저자는 언어 모델의 어텐션 레이어를 수정하는 것이 매우 효율적으로 죻은 성능을 얻기에 충분하다는 것을 발견했습니다. 이것이 LoRA 가중치를 모델의 어텐션 레이어에 추가하는 것이 일반적인 이유입니다. LoRA 작동 방식에 대한 자세한 내용은 [Using LoRA for effective Stable Diffusion fine-tuning](https://huggingface.co/blog/lora) 블로그를 확인하세요! [cloneofsimo](https://github.com/cloneofsimo)는 인기 있는 [lora](https://github.com/cloneofsimo/lora) GitHub 리포지토리에서 Stable Diffusion을 위한 LoRA 학습을 최초로 시도했습니다. 🧨 Diffusers는 [text-to-image 생성](https://github.com/huggingface/diffusers/tree/main/examples/text_to_image#training-with-lora) 및 [DreamBooth](https://github.com/huggingface/diffusers/tree/main/examples/dreambooth#training-with-low-rank-adaptation-of-large-language-models-lora)을 지원합니다. 이 가이드는 두 가지를 모두 수행하는 방법을 보여줍니다. @@ -104,11 +98,8 @@ accelerate launch train_dreambooth_lora.py \ *기본 모델의 가중치 위에* 파인튜닝된 DreamBooth 모델에서 LoRA 가중치를 불러온 다음, 더 빠른 추론을 위해 파이프라인을 GPU로 이동합니다. LoRA 가중치를 프리징된 사전 훈련된 모델 가중치와 병합할 때, 선택적으로 'scale' 매개변수로 어느 정도의 가중치를 병합할 지 조절할 수 있습니다: - - -💡 `0`의 `scale` 값은 LoRA 가중치를 사용하지 않아 원래 모델의 가중치만 사용한 것과 같고, `1`의 `scale` 값은 파인튜닝된 LoRA 가중치만 사용함을 의미합니다. 0과 1 사이의 값들은 두 결과들 사이로 보간됩니다. - - +> [!TIP] +> 💡 `0`의 `scale` 값은 LoRA 가중치를 사용하지 않아 원래 모델의 가중치만 사용한 것과 같고, `1`의 `scale` 값은 파인튜닝된 LoRA 가중치만 사용함을 의미합니다. 0과 1 사이의 값들은 두 결과들 사이로 보간됩니다. ```py >>> pipe.unet.load_attn_procs(model_path) diff --git a/docs/source/ko/training/text2image.md b/docs/source/ko/training/text2image.md index 4283f73ed9..b26603bf1b 100644 --- a/docs/source/ko/training/text2image.md +++ b/docs/source/ko/training/text2image.md @@ -13,11 +13,8 @@ specific language governing permissions and limitations under the License. # Text-to-image - - -text-to-image 파인튜닝 스크립트는 experimental 상태입니다. 과적합하기 쉽고 치명적인 망각과 같은 문제에 부딪히기 쉽습니다. 자체 데이터셋에서 최상의 결과를 얻으려면 다양한 하이퍼파라미터를 탐색하는 것이 좋습니다. - - +> [!WARNING] +> text-to-image 파인튜닝 스크립트는 experimental 상태입니다. 과적합하기 쉽고 치명적인 망각과 같은 문제에 부딪히기 쉽습니다. 자체 데이터셋에서 최상의 결과를 얻으려면 다양한 하이퍼파라미터를 탐색하는 것이 좋습니다. Stable Diffusion과 같은 text-to-image 모델은 텍스트 프롬프트에서 이미지를 생성합니다. 이 가이드는 PyTorch 및 Flax를 사용하여 자체 데이터셋에서 [`CompVis/stable-diffusion-v1-4`](https://huggingface.co/CompVis/stable-diffusion-v1-4) 모델로 파인튜닝하는 방법을 보여줍니다. 이 가이드에 사용된 text-to-image 파인튜닝을 위한 모든 학습 스크립트에 관심이 있는 경우 이 [리포지토리](https://github.com/huggingface/diffusers/tree/main/examples/text_to_image)에서 자세히 찾을 수 있습니다. diff --git a/docs/source/ko/training/text_inversion.md b/docs/source/ko/training/text_inversion.md index b27bed7d14..d8b44930e3 100644 --- a/docs/source/ko/training/text_inversion.md +++ b/docs/source/ko/training/text_inversion.md @@ -23,11 +23,8 @@ specific language governing permissions and limitations under the License. 이 가이드에서는 textual-inversion으로 [`stable-diffusion-v1-5/stable-diffusion-v1-5`](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) 모델을 학습하는 방법을 설명합니다. 이 가이드에서 사용된 모든 textual-inversion 학습 스크립트는 [여기](https://github.com/huggingface/diffusers/tree/main/examples/textual_inversion)에서 확인할 수 있습니다. 내부적으로 어떻게 작동하는지 자세히 살펴보고 싶으시다면 해당 링크를 참조해주시기 바랍니다. - - -[Stable Diffusion Textual Inversion Concepts Library](https://huggingface.co/sd-concepts-library)에는 커뮤니티에서 제작한 학습된 textual-inversion 모델들이 있습니다. 시간이 지남에 따라 더 많은 콘셉트들이 추가되어 유용한 리소스로 성장할 것입니다! - - +> [!TIP] +> [Stable Diffusion Textual Inversion Concepts Library](https://huggingface.co/sd-concepts-library)에는 커뮤니티에서 제작한 학습된 textual-inversion 모델들이 있습니다. 시간이 지남에 따라 더 많은 콘셉트들이 추가되어 유용한 리소스로 성장할 것입니다! 시작하기 전에 학습을 위한 의존성 라이브러리들을 설치해야 합니다: @@ -100,11 +97,8 @@ snapshot_download( - `token_identifier.txt` - `type_of_concept.txt`. - - -💡V100 GPU 1개를 기준으로 전체 학습에는 최대 1시간이 걸립니다. 학습이 완료되기를 기다리는 동안 궁금한 점이 있으면 아래 섹션에서 [textual-inversion이 어떻게 작동하는지](https://huggingface.co/docs/diffusers/training/text_inversion#how-it-works) 자유롭게 확인하세요 ! - - +> [!TIP] +> 💡V100 GPU 1개를 기준으로 전체 학습에는 최대 1시간이 걸립니다. 학습이 완료되기를 기다리는 동안 궁금한 점이 있으면 아래 섹션에서 [textual-inversion이 어떻게 작동하는지](https://huggingface.co/docs/diffusers/training/text_inversion#how-it-works) 자유롭게 확인하세요 ! @@ -128,15 +122,12 @@ accelerate launch textual_inversion.py \ --push_to_hub ``` - - -💡학습 성능을 올리기 위해, 플레이스홀더 토큰(``)을 (단일한 임베딩 벡터가 아닌) 복수의 임베딩 벡터로 표현하는 것 역시 고려할 있습니다. 이러한 트릭이 모델이 보다 복잡한 이미지의 스타일(앞서 말한 콘셉트)을 더 잘 캡처하는 데 도움이 될 수 있습니다. 복수의 임베딩 벡터 학습을 활성화하려면 다음 옵션을 전달하십시오. - -```bash ---num_vectors=5 -``` - - +> [!TIP] +> 💡학습 성능을 올리기 위해, 플레이스홀더 토큰(``)을 (단일한 임베딩 벡터가 아닌) 복수의 임베딩 벡터로 표현하는 것 역시 고려할 있습니다. 이러한 트릭이 모델이 보다 복잡한 이미지의 스타일(앞서 말한 콘셉트)을 더 잘 캡처하는 데 도움이 될 수 있습니다. 복수의 임베딩 벡터 학습을 활성화하려면 다음 옵션을 전달하십시오. +> +> ```bash +> --num_vectors=5 +> ``` @@ -193,11 +184,8 @@ textual-inversion 스크립트는 기본적으로 textual-inversion을 통해 - - -💡 커뮤니티는 [sd-concepts-library](https://huggingface.co/sd-concepts-library) 라는 대규모의 textual-inversion 임베딩 벡터 라이브러리를 만들었습니다. textual-inversion 임베딩을 밑바닥부터 학습하는 대신, 해당 라이브러리에 본인이 찾는 textual-inversion 임베딩이 이미 추가되어 있지 않은지를 확인하는 것도 좋은 방법이 될 것 같습니다. - - +> [!TIP] +> 💡 커뮤니티는 [sd-concepts-library](https://huggingface.co/sd-concepts-library) 라는 대규모의 textual-inversion 임베딩 벡터 라이브러리를 만들었습니다. textual-inversion 임베딩을 밑바닥부터 학습하는 대신, 해당 라이브러리에 본인이 찾는 textual-inversion 임베딩이 이미 추가되어 있지 않은지를 확인하는 것도 좋은 방법이 될 것 같습니다. textual-inversion 임베딩 벡터을 불러오기 위해서는, 먼저 해당 임베딩 벡터를 학습할 때 사용한 모델을 불러와야 합니다. 여기서는 [`stable-diffusion-v1-5/stable-diffusion-v1-5`](https://huggingface.co/docs/diffusers/training/stable-diffusion-v1-5/stable-diffusion-v1-5) 모델이 사용되었다고 가정하고 불러오겠습니다. diff --git a/docs/source/ko/training/unconditional_training.md b/docs/source/ko/training/unconditional_training.md index c8c463da6b..04a9a6c7ea 100644 --- a/docs/source/ko/training/unconditional_training.md +++ b/docs/source/ko/training/unconditional_training.md @@ -78,11 +78,8 @@ write_basic_config() 학습 스크립트는 `diffusion_pytorch_model.bin` 파일을 생성하고, 그것을 당신의 리포지토리에 저장합니다. - - -💡 전체 학습은 V100 GPU 4개를 사용할 경우, 2시간이 소요됩니다. - - +> [!TIP] +> 💡 전체 학습은 V100 GPU 4개를 사용할 경우, 2시간이 소요됩니다. 예를 들어, [Oxford Flowers](https://huggingface.co/datasets/huggan/flowers-102-categories) 데이터셋을 사용해 파인튜닝할 경우: diff --git a/docs/source/ko/tutorials/basic_training.md b/docs/source/ko/tutorials/basic_training.md index 2c4c89edd1..05ce1037b5 100644 --- a/docs/source/ko/tutorials/basic_training.md +++ b/docs/source/ko/tutorials/basic_training.md @@ -19,11 +19,8 @@ Unconditional 이미지 생성은 학습에 사용된 데이터셋과 유사한 이 튜토리얼은 나만의 🦋 나비 🦋를 생성하기 위해 [Smithsonian Butterflies](https://huggingface.co/datasets/huggan/smithsonian_butterflies_subset) 데이터셋의 하위 집합에서 [`UNet2DModel`] 모델을 학습하는 방법을 가르쳐줄 것입니다. - - -💡 이 학습 튜토리얼은 [Training with 🧨 Diffusers](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/training_example.ipynb) 노트북 기반으로 합니다. Diffusion 모델의 작동 방식 및 자세한 내용은 노트북을 확인하세요! - - +> [!TIP] +> 💡 이 학습 튜토리얼은 [Training with 🧨 Diffusers](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/training_example.ipynb) 노트북 기반으로 합니다. Diffusion 모델의 작동 방식 및 자세한 내용은 노트북을 확인하세요! 시작 전에, 🤗 Datasets을 불러오고 전처리하기 위해 데이터셋이 설치되어 있는지 다수 GPU에서 학습을 간소화하기 위해 🤗 Accelerate 가 설치되어 있는지 확인하세요. 그 후 학습 메트릭을 시각화하기 위해 [TensorBoard](https://www.tensorflow.org/tensorboard)를 또한 설치하세요. (또한 학습 추적을 위해 [Weights & Biases](https://docs.wandb.ai/)를 사용할 수 있습니다.) diff --git a/docs/source/ko/using-diffusers/controlling_generation.md b/docs/source/ko/using-diffusers/controlling_generation.md index 1b9a8b5df5..db22fe042d 100644 --- a/docs/source/ko/using-diffusers/controlling_generation.md +++ b/docs/source/ko/using-diffusers/controlling_generation.md @@ -85,12 +85,9 @@ Pix2Pix Zero는 합성 이미지와 실제 이미지를 편집하는 데 모두 다음으로 편집할 컨셉과 새로운 타겟 컨셉에 대한 이미지 캡션을 생성합니다. 이를 위해 [Flan-T5](https://huggingface.co/docs/transformers/model_doc/flan-t5)와 같은 모델을 사용할 수 있습니다. 그런 다음 텍스트 인코더를 통해 소스 개념과 대상 개념 모두에 대한 "평균" 프롬프트 임베딩을 생성합니다. 마지막으로, 합성 이미지를 편집하기 위해 pix2pix-zero 알고리즘을 사용합니다. - 실제 이미지를 편집하려면 먼저 [BLIP](https://huggingface.co/docs/transformers/model_doc/blip)과 같은 모델을 사용하여 이미지 캡션을 생성합니다. 그런 다음 프롬프트와 이미지에 ddim 반전을 적용하여 "역(inverse)" latents을 생성합니다. 이전과 마찬가지로 소스 및 대상 개념 모두에 대한 "평균(mean)" 프롬프트 임베딩이 생성되고 마지막으로 "역(inverse)" latents와 결합된 pix2pix-zero 알고리즘이 이미지를 편집하는 데 사용됩니다. - - -Pix2Pix Zero는 '제로 샷(zero-shot)' 이미지 편집이 가능한 최초의 모델입니다. -즉, 이 모델은 다음과 같이 일반 소비자용 GPU에서 1분 이내에 이미지를 편집할 수 있습니다(../api/pipelines/stable_diffusion/pix2pix_zero#usage-example). - - +> [!TIP] +> Pix2Pix Zero는 '제로 샷(zero-shot)' 이미지 편집이 가능한 최초의 모델입니다. +> 즉, 이 모델은 다음과 같이 일반 소비자용 GPU에서 1분 이내에 이미지를 편집할 수 있습니다(../api/pipelines/stable_diffusion/pix2pix_zero#usage-example). 위에서 언급했듯이 Pix2Pix Zero에는 특정 개념으로 세대를 유도하기 위해 (UNet, VAE 또는 텍스트 인코더가 아닌) latents을 최적화하는 기능이 포함되어 있습니다.즉, 전체 파이프라인에 표준 [StableDiffusionPipeline](../api/pipelines/stable_diffusion/text2img)보다 더 많은 메모리가 필요할 수 있습니다. @@ -140,13 +137,10 @@ SAG는 고빈도 세부 정보를 기반으로 하지 않은 예측에서 완전 사용 방법에 대한 자세한 내용은 [여기](../api/pipelines/stable_diffusion_2#depthtoimage)를 참조하세요. - - -InstructPix2Pix와 Pix2Pix Zero와 같은 방법의 중요한 차이점은 전자의 경우 -는 사전 학습된 가중치를 미세 조정하는 반면, 후자는 그렇지 않다는 것입니다. 즉, 다음을 수행할 수 있습니다. -사용 가능한 모든 안정적 확산 모델에 Pix2Pix Zero를 적용할 수 있습니다. - - +> [!TIP] +> InstructPix2Pix와 Pix2Pix Zero와 같은 방법의 중요한 차이점은 전자의 경우 +> 는 사전 학습된 가중치를 미세 조정하는 반면, 후자는 그렇지 않다는 것입니다. 즉, 다음을 수행할 수 있습니다. +> 사용 가능한 모든 안정적 확산 모델에 Pix2Pix Zero를 적용할 수 있습니다. ## MultiDiffusion Panorama diff --git a/docs/source/ko/using-diffusers/custom_pipeline_overview.md b/docs/source/ko/using-diffusers/custom_pipeline_overview.md index b143bf8ab0..caeeca8cef 100644 --- a/docs/source/ko/using-diffusers/custom_pipeline_overview.md +++ b/docs/source/ko/using-diffusers/custom_pipeline_overview.md @@ -20,11 +20,8 @@ specific language governing permissions and limitations under the License. 허브에서 커뮤니티 파이프라인을 로드하려면, 커뮤니티 파이프라인의 리포지토리 ID와 (파이프라인 가중치 및 구성 요소를 로드하려는) 모델의 리포지토리 ID를 인자로 전달해야 합니다. 예를 들어, 아래 예시에서는 `hf-internal-testing/diffusers-dummy-pipeline`에서 더미 파이프라인을 불러오고, `google/ddpm-cifar10-32`에서 파이프라인의 가중치와 컴포넌트들을 로드합니다. - - -🔒 허깅 페이스 허브에서 커뮤니티 파이프라인을 불러오는 것은 곧 해당 코드가 안전하다고 신뢰하는 것입니다. 코드를 자동으로 불러오고 실행하기 앞서 반드시 온라인으로 해당 코드의 신뢰성을 검사하세요! - - +> [!WARNING] +> 🔒 허깅 페이스 허브에서 커뮤니티 파이프라인을 불러오는 것은 곧 해당 코드가 안전하다고 신뢰하는 것입니다. 코드를 자동으로 불러오고 실행하기 앞서 반드시 온라인으로 해당 코드의 신뢰성을 검사하세요! ```py from diffusers import DiffusionPipeline diff --git a/docs/source/ko/using-diffusers/diffedit.md b/docs/source/ko/using-diffusers/diffedit.md index 74b9e97831..edf23f0214 100644 --- a/docs/source/ko/using-diffusers/diffedit.md +++ b/docs/source/ko/using-diffusers/diffedit.md @@ -156,11 +156,8 @@ print(source_prompts) print(target_prompts) ``` - - -다양한 품질의 텍스트를 생성하는 전략에 대해 자세히 알아보려면 [생성 전략](https://huggingface.co/docs/transformers/main/en/generation_strategies) 가이드를 참조하세요. - - +> [!TIP] +> 다양한 품질의 텍스트를 생성하는 전략에 대해 자세히 알아보려면 [생성 전략](https://huggingface.co/docs/transformers/main/en/generation_strategies) 가이드를 참조하세요. 텍스트 인코딩을 위해 [`StableDiffusionDiffEditPipeline`]에서 사용하는 텍스트 인코더 모델을 불러옵니다. 텍스트 인코더를 사용하여 텍스트 임베딩을 계산합니다: diff --git a/docs/source/ko/using-diffusers/img2img.md b/docs/source/ko/using-diffusers/img2img.md index 8da840f748..3901fb755f 100644 --- a/docs/source/ko/using-diffusers/img2img.md +++ b/docs/source/ko/using-diffusers/img2img.md @@ -53,11 +53,8 @@ init_image - - -💡 `strength`는 입력 이미지에 추가되는 노이즈의 양을 제어하는 0.0에서 1.0 사이의 값입니다. 1.0에 가까운 값은 다양한 변형을 허용하지만 입력 이미지와 의미적으로 일치하지 않는 이미지를 생성합니다. - - +> [!TIP] +> 💡 `strength`는 입력 이미지에 추가되는 노이즈의 양을 제어하는 0.0에서 1.0 사이의 값입니다. 1.0에 가까운 값은 다양한 변형을 허용하지만 입력 이미지와 의미적으로 일치하지 않는 이미지를 생성합니다. 프롬프트를 정의하고(지브리 스타일(Ghibli-style)에 맞게 조정된 이 체크포인트의 경우 프롬프트 앞에 `ghibli style` 토큰을 붙여야 합니다) 파이프라인을 실행합니다: diff --git a/docs/source/ko/using-diffusers/inpaint.md b/docs/source/ko/using-diffusers/inpaint.md index adf1251176..cefb892186 100644 --- a/docs/source/ko/using-diffusers/inpaint.md +++ b/docs/source/ko/using-diffusers/inpaint.md @@ -59,11 +59,8 @@ image = pipe(prompt=prompt, image=init_image, mask_image=mask_image).images[0] :-------------------------:|:-------------------------:|:-------------------------:|-------------------------:| drawing | drawing | ***Face of a yellow cat, high resolution, sitting on a park bench*** | drawing | - - -이전의 실험적인 인페인팅 구현에서는 품질이 낮은 다른 프로세스를 사용했습니다. 이전 버전과의 호환성을 보장하기 위해 새 모델이 포함되지 않은 사전학습된 파이프라인을 불러오면 이전 인페인팅 방법이 계속 적용됩니다. - - +> [!WARNING] +> 이전의 실험적인 인페인팅 구현에서는 품질이 낮은 다른 프로세스를 사용했습니다. 이전 버전과의 호환성을 보장하기 위해 새 모델이 포함되지 않은 사전학습된 파이프라인을 불러오면 이전 인페인팅 방법이 계속 적용됩니다. 아래 Space에서 이미지 인페인팅을 직접 해보세요! diff --git a/docs/source/ko/using-diffusers/kandinsky.md b/docs/source/ko/using-diffusers/kandinsky.md index cc554c67f9..8eff8f5629 100644 --- a/docs/source/ko/using-diffusers/kandinsky.md +++ b/docs/source/ko/using-diffusers/kandinsky.md @@ -31,15 +31,12 @@ Kandinsky 모델은 일련의 다국어 text-to-image 생성 모델입니다. Ka #!pip install -q diffusers transformers accelerate ``` - - -Kandinsky 2.1과 2.2의 사용법은 매우 유사합니다! 유일한 차이점은 Kandinsky 2.2는 latents를 디코딩할 때 `프롬프트`를 입력으로 받지 않는다는 것입니다. 대신, Kandinsky 2.2는 디코딩 중에는 `image_embeds`만 받아들입니다. - -
- -Kandinsky 3는 더 간결한 아키텍처를 가지고 있으며 prior 모델이 필요하지 않습니다. 즉, [Stable Diffusion XL](sdxl)과 같은 다른 diffusion 모델과 사용법이 동일합니다. - - +> [!WARNING] +> Kandinsky 2.1과 2.2의 사용법은 매우 유사합니다! 유일한 차이점은 Kandinsky 2.2는 latents를 디코딩할 때 `프롬프트`를 입력으로 받지 않는다는 것입니다. 대신, Kandinsky 2.2는 디코딩 중에는 `image_embeds`만 받아들입니다. +> +>
+> +> Kandinsky 3는 더 간결한 아키텍처를 가지고 있으며 prior 모델이 필요하지 않습니다. 즉, [Stable Diffusion XL](sdxl)과 같은 다른 diffusion 모델과 사용법이 동일합니다. ## Text-to-image @@ -321,20 +318,17 @@ make_image_grid([original_image.resize((512, 512)), image.resize((512, 512))], r ## Inpainting - - -⚠️ Kandinsky 모델은 이제 검은색 픽셀 대신 ⬜️ **흰색 픽셀**을 사용하여 마스크 영역을 표현합니다. 프로덕션에서 [`KandinskyInpaintPipeline`]을 사용하는 경우 흰색 픽셀을 사용하도록 마스크를 변경해야 합니다: - -```py -# PIL 입력에 대해 -import PIL.ImageOps -mask = PIL.ImageOps.invert(mask) - -# PyTorch와 NumPy 입력에 대해 -mask = 1 - mask -``` - - +> [!WARNING] +> ⚠️ Kandinsky 모델은 이제 검은색 픽셀 대신 ⬜️ **흰색 픽셀**을 사용하여 마스크 영역을 표현합니다. 프로덕션에서 [`KandinskyInpaintPipeline`]을 사용하는 경우 흰색 픽셀을 사용하도록 마스크를 변경해야 합니다: +> +> ```py +> # PIL 입력에 대해 +> import PIL.ImageOps +> mask = PIL.ImageOps.invert(mask) +> +> # PyTorch와 NumPy 입력에 대해 +> mask = 1 - mask +> ``` 인페인팅에서는 원본 이미지, 원본 이미지에서 대체할 영역의 마스크, 인페인팅할 내용에 대한 텍스트 프롬프트가 필요합니다. Prior 파이프라인을 불러옵니다: @@ -565,11 +559,8 @@ image ## ControlNet - - -⚠️ ControlNet은 Kandinsky 2.2에서만 지원됩니다! - - +> [!WARNING] +> ⚠️ ControlNet은 Kandinsky 2.2에서만 지원됩니다! ControlNet을 사용하면 depth map이나 edge detection와 같은 추가 입력을 통해 사전학습된 large diffusion 모델을 conditioning할 수 있습니다. 예를 들어, 모델이 depth map의 구조를 이해하고 보존할 수 있도록 깊이 맵으로 Kandinsky 2.2를 conditioning할 수 있습니다. diff --git a/docs/source/ko/using-diffusers/loading.md b/docs/source/ko/using-diffusers/loading.md index 3d6b7634b4..2160acacc2 100644 --- a/docs/source/ko/using-diffusers/loading.md +++ b/docs/source/ko/using-diffusers/loading.md @@ -30,11 +30,8 @@ diffusion 모델의 훈련과 추론에 필요한 모든 것은 [`DiffusionPipel ## Diffusion 파이프라인 - - -💡 [`DiffusionPipeline`] 클래스가 동작하는 방식에 보다 자세한 내용이 궁금하다면, [DiffusionPipeline explained](#diffusionpipeline에-대해-알아보기) 섹션을 확인해보세요. - - +> [!TIP] +> 💡 [`DiffusionPipeline`] 클래스가 동작하는 방식에 보다 자세한 내용이 궁금하다면, [DiffusionPipeline explained](#diffusionpipeline에-대해-알아보기) 섹션을 확인해보세요. [`DiffusionPipeline`] 클래스는 diffusion 모델을 [허브](https://huggingface.co/models?library=diffusers)로부터 불러오는 가장 심플하면서 보편적인 방식입니다. [`DiffusionPipeline.from_pretrained`] 메서드는 적합한 파이프라인 클래스를 자동으로 탐지하고, 필요한 구성요소(configuration)와 가중치(weight) 파일들을 다운로드하고 캐싱한 다음, 해당 파이프라인 인스턴스를 반환합니다. @@ -175,11 +172,8 @@ Variant란 일반적으로 다음과 같은 체크포인트들을 의미합니 - `torch.float16`과 같이 정밀도는 더 낮지만, 용량 역시 더 작은 부동소수점 타입의 가중치를 사용하는 체크포인트. *(다만 이와 같은 variant의 경우, 추가적인 훈련과 CPU환경에서의 구동이 불가능합니다.)* - Non-EMA 가중치를 사용하는 체크포인트. *(Non-EMA 가중치의 경우, 파인 튜닝 단계에서 사용하는 것이 권장되는데, 추론 단계에선 사용하지 않는 것이 권장됩니다.)* - - -💡 모델 구조는 동일하지만 서로 다른 학습 환경에서 서로 다른 데이터셋으로 학습된 체크포인트들이 있을 경우, 해당 체크포인트들은 variant 단계가 아닌 리포지토리 단계에서 분리되어 관리되어야 합니다. (즉, 해당 체크포인트들은 서로 다른 리포지토리에서 따로 관리되어야 합니다. 예시: [`stable-diffusion-v1-4`], [`stable-diffusion-v1-5`]). - - +> [!TIP] +> 💡 모델 구조는 동일하지만 서로 다른 학습 환경에서 서로 다른 데이터셋으로 학습된 체크포인트들이 있을 경우, 해당 체크포인트들은 variant 단계가 아닌 리포지토리 단계에서 분리되어 관리되어야 합니다. (즉, 해당 체크포인트들은 서로 다른 리포지토리에서 따로 관리되어야 합니다. 예시: [`stable-diffusion-v1-4`], [`stable-diffusion-v1-5`]). | **checkpoint type** | **weight name** | **argument for loading weights** | | ------------------- | ----------------------------------- | -------------------------------- | diff --git a/docs/source/ko/using-diffusers/loading_adapters.md b/docs/source/ko/using-diffusers/loading_adapters.md index f0d085bc6a..e7ae116575 100644 --- a/docs/source/ko/using-diffusers/loading_adapters.md +++ b/docs/source/ko/using-diffusers/loading_adapters.md @@ -18,11 +18,8 @@ specific language governing permissions and limitations under the License. 이 가이드에서는 DreamBooth, textual inversion 및 LoRA 가중치를 불러오는 방법을 설명합니다. - - -사용할 체크포인트와 임베딩은 [Stable Diffusion Conceptualizer](https://huggingface.co/spaces/sd-concepts-library/stable-diffusion-conceptualizer), [LoRA the Explorer](https://huggingface.co/spaces/multimodalart/LoraTheExplorer), [Diffusers Models Gallery](https://huggingface.co/spaces/huggingface-projects/diffusers-gallery)에서 찾아보시기 바랍니다. - - +> [!TIP] +> 사용할 체크포인트와 임베딩은 [Stable Diffusion Conceptualizer](https://huggingface.co/spaces/sd-concepts-library/stable-diffusion-conceptualizer), [LoRA the Explorer](https://huggingface.co/spaces/multimodalart/LoraTheExplorer), [Diffusers Models Gallery](https://huggingface.co/spaces/huggingface-projects/diffusers-gallery)에서 찾아보시기 바랍니다. ## DreamBooth @@ -101,11 +98,8 @@ image [Low-Rank Adaptation (LoRA)](https://huggingface.co/papers/2106.09685)은 속도가 빠르고 파일 크기가 (수백 MB로) 작기 때문에 널리 사용되는 학습 기법입니다. 이 가이드의 다른 방법과 마찬가지로, LoRA는 몇 장의 이미지만으로 새로운 스타일을 학습하도록 모델을 학습시킬 수 있습니다. 이는 diffusion 모델에 새로운 가중치를 삽입한 다음 전체 모델 대신 새로운 가중치만 학습시키는 방식으로 작동합니다. 따라서 LoRA를 더 빠르게 학습시키고 더 쉽게 저장할 수 있습니다. - - -LoRA는 다른 학습 방법과 함께 사용할 수 있는 매우 일반적인 학습 기법입니다. 예를 들어, DreamBooth와 LoRA로 모델을 학습하는 것이 일반적입니다. 또한 새롭고 고유한 이미지를 생성하기 위해 여러 개의 LoRA를 불러오고 병합하는 것이 점점 더 일반화되고 있습니다. 병합은 이 불러오기 가이드의 범위를 벗어나므로 자세한 내용은 심층적인 [LoRA 병합](merge_loras) 가이드에서 확인할 수 있습니다. - - +> [!TIP] +> LoRA는 다른 학습 방법과 함께 사용할 수 있는 매우 일반적인 학습 기법입니다. 예를 들어, DreamBooth와 LoRA로 모델을 학습하는 것이 일반적입니다. 또한 새롭고 고유한 이미지를 생성하기 위해 여러 개의 LoRA를 불러오고 병합하는 것이 점점 더 일반화되고 있습니다. 병합은 이 불러오기 가이드의 범위를 벗어나므로 자세한 내용은 심층적인 [LoRA 병합](merge_loras) 가이드에서 확인할 수 있습니다. LoRA는 다른 모델과 함께 사용해야 합니다: @@ -184,11 +178,8 @@ pipe.set_adapters("my_adapter", scales) 이는 여러 어댑터에서도 작동합니다. 방법은 [이 가이드](https://huggingface.co/docs/diffusers/tutorials/using_peft_for_inference#customize-adapters-strength)를 참조하세요. - - -현재 [`~loaders.LoraLoaderMixin.set_adapters`]는 어텐션 가중치의 스케일링만 지원합니다. LoRA에 다른 부분(예: resnets or down-/upsamplers)이 있는 경우 1.0의 스케일을 유지합니다. - - +> [!WARNING] +> 현재 [`~loaders.LoraLoaderMixin.set_adapters`]는 어텐션 가중치의 스케일링만 지원합니다. LoRA에 다른 부분(예: resnets or down-/upsamplers)이 있는 경우 1.0의 스케일을 유지합니다. ### Kohya와 TheLastBen @@ -222,14 +213,11 @@ image = pipeline(prompt).images[0] image ``` - - -Kohya LoRA를 🤗 Diffusers와 함께 사용할 때 몇 가지 제한 사항이 있습니다: - -- [여기](https://github.com/huggingface/diffusers/pull/4287/#issuecomment-1655110736)에 설명된 여러 가지 이유로 인해 이미지가 ComfyUI와 같은 UI에서 생성된 이미지와 다르게 보일 수 있습니다. -- [LyCORIS 체크포인트](https://github.com/KohakuBlueleaf/LyCORIS)가 완전히 지원되지 않습니다. [`~loaders.LoraLoaderMixin.load_lora_weights`] 메서드는 LoRA 및 LoCon 모듈로 LyCORIS 체크포인트를 불러올 수 있지만, Hada 및 LoKR은 지원되지 않습니다. - - +> [!WARNING] +> Kohya LoRA를 🤗 Diffusers와 함께 사용할 때 몇 가지 제한 사항이 있습니다: +> +> - [여기](https://github.com/huggingface/diffusers/pull/4287/#issuecomment-1655110736)에 설명된 여러 가지 이유로 인해 이미지가 ComfyUI와 같은 UI에서 생성된 이미지와 다르게 보일 수 있습니다. +> - [LyCORIS 체크포인트](https://github.com/KohakuBlueleaf/LyCORIS)가 완전히 지원되지 않습니다. [`~loaders.LoraLoaderMixin.load_lora_weights`] 메서드는 LoRA 및 LoCon 모듈로 LyCORIS 체크포인트를 불러올 수 있지만, Hada 및 LoKR은 지원되지 않습니다. @@ -326,9 +314,8 @@ pipeline.load_ip_adapter("h94/IP-Adapter", subfolder="sdxl_models", weight_name= IP-Adapter FaceID 모델은 CLIP 이미지 임베딩 대신 `insightface`에서 생성한 이미지 임베딩을 사용하는 실험적인 IP Adapter입니다. 이러한 모델 중 일부는 LoRA를 사용하여 ID 일관성을 개선하기도 합니다. 이러한 모델을 사용하려면 `insightface`와 해당 요구 사항을 모두 설치해야 합니다. - -InsightFace 사전학습된 모델은 비상업적 연구 목적으로만 사용할 수 있으므로, IP-Adapter-FaceID 모델은 연구 목적으로만 릴리즈되었으며 상업적 용도로는 사용할 수 없습니다. - +> [!WARNING] +> InsightFace 사전학습된 모델은 비상업적 연구 목적으로만 사용할 수 있으므로, IP-Adapter-FaceID 모델은 연구 목적으로만 릴리즈되었으며 상업적 용도로는 사용할 수 없습니다. ```py pipeline = AutoPipelineForText2Image.from_pretrained( diff --git a/docs/source/ko/using-diffusers/other-formats.md b/docs/source/ko/using-diffusers/other-formats.md index 3034551f48..f5a71f56eb 100644 --- a/docs/source/ko/using-diffusers/other-formats.md +++ b/docs/source/ko/using-diffusers/other-formats.md @@ -14,11 +14,8 @@ specific language governing permissions and limitations under the License. Stable Diffusion 모델들은 학습 및 저장된 프레임워크와 다운로드 위치에 따라 다양한 형식으로 제공됩니다. 이러한 형식을 🤗 Diffusers에서 사용할 수 있도록 변환하면 추론을 위한 [다양한 스케줄러 사용](schedulers), 사용자 지정 파이프라인 구축, 추론 속도 최적화를 위한 다양한 기법과 방법 등 라이브러리에서 지원하는 모든 기능을 사용할 수 있습니다. - - -우리는 `.safetensors` 형식을 추천합니다. 왜냐하면 기존의 pickled 파일은 취약하고 머신에서 코드를 실행할 때 악용될 수 있는 것에 비해 훨씬 더 안전합니다. (safetensors 불러오기 가이드에서 자세히 알아보세요.) - - +> [!TIP] +> 우리는 `.safetensors` 형식을 추천합니다. 왜냐하면 기존의 pickled 파일은 취약하고 머신에서 코드를 실행할 때 악용될 수 있는 것에 비해 훨씬 더 안전합니다. (safetensors 불러오기 가이드에서 자세히 알아보세요.) 이 가이드에서는 다른 Stable Diffusion 형식을 🤗 Diffusers와 호환되도록 변환하는 방법을 설명합니다. diff --git a/docs/source/ko/using-diffusers/schedulers.md b/docs/source/ko/using-diffusers/schedulers.md index 55424c9982..b12c08b8c8 100644 --- a/docs/source/ko/using-diffusers/schedulers.md +++ b/docs/source/ko/using-diffusers/schedulers.md @@ -318,12 +318,9 @@ images = pipeline(prompt_ids, params, prng_seed, num_inference_steps, jit=True). images = pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]))) ``` - - -다음 Flax 스케줄러는 *아직* Flax Stable Diffusion 파이프라인과 호환되지 않습니다. - -- `FlaxLMSDiscreteScheduler` -- `FlaxDDPMScheduler` - - +> [!WARNING] +> 다음 Flax 스케줄러는 *아직* Flax Stable Diffusion 파이프라인과 호환되지 않습니다. +> +> - `FlaxLMSDiscreteScheduler` +> - `FlaxDDPMScheduler` diff --git a/docs/source/ko/using-diffusers/shap-e.md b/docs/source/ko/using-diffusers/shap-e.md index abf5a182b3..4c9d7fb7d1 100644 --- a/docs/source/ko/using-diffusers/shap-e.md +++ b/docs/source/ko/using-diffusers/shap-e.md @@ -151,11 +151,8 @@ images = pipe(prompt, guidance_scale=guidance_scale, num_inference_steps=64, fra 메시 출력을 `ply` 파일로 저장하려면 [`~utils.export_to_ply`] 함수를 사용합니다: - - -선택적으로 [`~utils.export_to_obj`] 함수를 사용하여 메시 출력을 `obj` 파일로 저장할 수 있습니다. 다양한 형식으로 메시 출력을 저장할 수 있어 다운스트림에서 더욱 유연하게 사용할 수 있습니다! - - +> [!TIP] +> 선택적으로 [`~utils.export_to_obj`] 함수를 사용하여 메시 출력을 `obj` 파일로 저장할 수 있습니다. 다양한 형식으로 메시 출력을 저장할 수 있어 다운스트림에서 더욱 유연하게 사용할 수 있습니다! ```py from diffusers.utils import export_to_ply diff --git a/docs/source/ko/using-diffusers/unconditional_image_generation.md b/docs/source/ko/using-diffusers/unconditional_image_generation.md index c3eaac4b03..b8fe800578 100644 --- a/docs/source/ko/using-diffusers/unconditional_image_generation.md +++ b/docs/source/ko/using-diffusers/unconditional_image_generation.md @@ -20,11 +20,8 @@ Unconditional 이미지 생성은 비교적 간단한 작업입니다. 모델이 먼저 ['DiffusionPipeline']의 인스턴스를 생성하고 다운로드할 파이프라인의 [체크포인트](https://huggingface.co/models?library=diffusers&sort=downloads)를 지정합니다. 허브의 🧨 diffusion 체크포인트 중 하나를 사용할 수 있습니다(사용할 체크포인트는 나비 이미지를 생성합니다). - - -💡 나만의 unconditional 이미지 생성 모델을 학습시키고 싶으신가요? 학습 가이드를 살펴보고 나만의 이미지를 생성하는 방법을 알아보세요. - - +> [!TIP] +> 💡 나만의 unconditional 이미지 생성 모델을 학습시키고 싶으신가요? 학습 가이드를 살펴보고 나만의 이미지를 생성하는 방법을 알아보세요. 이 가이드에서는 unconditional 이미지 생성에 ['DiffusionPipeline']과 [DDPM](https://huggingface.co/papers/2006.11239)을 사용합니다: diff --git a/docs/source/ko/using-diffusers/write_own_pipeline.md b/docs/source/ko/using-diffusers/write_own_pipeline.md index 45678763cc..ae6ce238ac 100644 --- a/docs/source/ko/using-diffusers/write_own_pipeline.md +++ b/docs/source/ko/using-diffusers/write_own_pipeline.md @@ -110,11 +110,8 @@ Stable Diffusion 은 text-to-image *latent diffusion* 모델입니다. latent di 보시다시피, 이것은 UNet 모델만 포함된 DDPM 파이프라인보다 더 복잡합니다. Stable Diffusion 모델에는 세 개의 개별 사전학습된 모델이 있습니다. - - -💡 VAE, UNet 및 텍스트 인코더 모델의 작동방식에 대한 자세한 내용은 [How does Stable Diffusion work?](https://huggingface.co/blog/stable_diffusion#how-does-stable-diffusion-work) 블로그를 참조하세요. - - +> [!TIP] +> 💡 VAE, UNet 및 텍스트 인코더 모델의 작동방식에 대한 자세한 내용은 [How does Stable Diffusion work?](https://huggingface.co/blog/stable_diffusion#how-does-stable-diffusion-work) 블로그를 참조하세요. 이제 Stable Diffusion 파이프라인에 필요한 구성요소들이 무엇인지 알았으니, [`~ModelMixin.from_pretrained`] 메서드를 사용해 모든 구성요소를 불러옵니다. 사전학습된 체크포인트 [`stable-diffusion-v1-5/stable-diffusion-v1-5`](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5)에서 찾을 수 있으며, 각 구성요소들은 별도의 하위 폴더에 저장되어 있습니다: @@ -151,11 +148,8 @@ Stable Diffusion 은 text-to-image *latent diffusion* 모델입니다. latent di 다음 단계는 임베딩을 생성하기 위해 텍스트를 토큰화하는 것입니다. 이 텍스트는 UNet 모델에서 condition으로 사용되고 입력 프롬프트와 유사한 방향으로 diffusion 프로세스를 조정하는 데 사용됩니다. - - -💡 `guidance_scale` 매개변수는 이미지를 생성할 때 프롬프트에 얼마나 많은 가중치를 부여할지 결정합니다. - - +> [!TIP] +> 💡 `guidance_scale` 매개변수는 이미지를 생성할 때 프롬프트에 얼마나 많은 가중치를 부여할지 결정합니다. 다른 프롬프트를 생성하고 싶다면 원하는 프롬프트를 자유롭게 선택하세요! @@ -198,15 +192,12 @@ Stable Diffusion 은 text-to-image *latent diffusion* 모델입니다. latent di 그다음 diffusion 프로세스의 시작점으로 초기 랜덤 노이즈를 생성합니다. 이것이 이미지의 잠재적 표현이며 점차적으로 노이즈가 제거됩니다. 이 시점에서 `latent` 이미지는 최종 이미지 크기보다 작지만 나중에 모델이 이를 512x512 이미지 크기로 변환하므로 괜찮습니다. - - -💡 `vae` 모델에는 3개의 다운 샘플링 레이어가 있기 때문에 높이와 너비가 8로 나뉩니다. 다음을 실행하여 확인할 수 있습니다: - -```py -2 ** (len(vae.config.block_out_channels) - 1) == 8 -``` - - +> [!TIP] +> 💡 `vae` 모델에는 3개의 다운 샘플링 레이어가 있기 때문에 높이와 너비가 8로 나뉩니다. 다음을 실행하여 확인할 수 있습니다: +> +> ```py +> 2 ** (len(vae.config.block_out_channels) - 1) == 8 +> ``` ```py >>> latents = torch.randn( diff --git a/docs/source/pt/installation.md b/docs/source/pt/installation.md index 1e83e36ca1..acc767110c 100644 --- a/docs/source/pt/installation.md +++ b/docs/source/pt/installation.md @@ -104,11 +104,8 @@ Esses comandos irá linkar a pasta que você clonou o repositório e os caminhos Python então irá procurar dentro da pasta que você clonou além dos caminhos normais das bibliotecas. Por exemplo, se o pacote python for tipicamente instalado no `~/anaconda3/envs/main/lib/python3.10/site-packages/`, o Python também irá procurar na pasta `~/diffusers/` que você clonou. - - -Você deve deixar a pasta `diffusers` se você quiser continuar usando a biblioteca. - - +> [!WARNING] +> Você deve deixar a pasta `diffusers` se você quiser continuar usando a biblioteca. Agora você pode facilmente atualizar seu clone para a última versão do 🤗 Diffusers com o seguinte comando: diff --git a/docs/source/pt/quicktour.md b/docs/source/pt/quicktour.md index 109f7e2712..5996b65a9c 100644 --- a/docs/source/pt/quicktour.md +++ b/docs/source/pt/quicktour.md @@ -24,11 +24,8 @@ Seja você um desenvolvedor ou um usuário, esse tour rápido irá introduzir vo Esse tour rápido mostrará como usar o [`DiffusionPipeline`] para inferência, e então mostrará como combinar um modelo e um agendador para replicar o que está acontecendo dentro do [`DiffusionPipeline`]. - - -Esse tour rápido é uma versão simplificada da introdução 🧨 Diffusers [notebook](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/diffusers_intro.ipynb) para ajudar você a começar rápido. Se você quer aprender mais sobre o objetivo do 🧨 Diffusers, filosofia de design, e detalhes adicionais sobre a API principal, veja o notebook! - - +> [!TIP] +> Esse tour rápido é uma versão simplificada da introdução 🧨 Diffusers [notebook](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/diffusers_intro.ipynb) para ajudar você a começar rápido. Se você quer aprender mais sobre o objetivo do 🧨 Diffusers, filosofia de design, e detalhes adicionais sobre a API principal, veja o notebook! Antes de começar, certifique-se de ter todas as bibliotecas necessárias instaladas: @@ -56,11 +53,8 @@ Comece criando uma instância do [`DiffusionPipeline`] e especifique qual checkp Você pode usar o [`DiffusionPipeline`] para qualquer [checkpoint](https://huggingface.co/models?library=diffusers&sort=downloads) armazenado no Hugging Face Hub. Nesse quicktour, você carregará o checkpoint [`stable-diffusion-v1-5`](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) para geração de texto para imagem. - - -Para os modelos de [Stable Diffusion](https://huggingface.co/CompVis/stable-diffusion), por favor leia cuidadosamente a [licença](https://huggingface.co/spaces/CompVis/stable-diffusion-license) primeiro antes de rodar o modelo. 🧨 Diffusers implementa uma verificação de segurança: [`safety_checker`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/safety_checker.py) para prevenir conteúdo ofensivo ou nocivo, mas as capacidades de geração de imagem aprimorada do modelo podem ainda produzir conteúdo potencialmente nocivo. - - +> [!WARNING] +> Para os modelos de [Stable Diffusion](https://huggingface.co/CompVis/stable-diffusion), por favor leia cuidadosamente a [licença](https://huggingface.co/spaces/CompVis/stable-diffusion-license) primeiro antes de rodar o modelo. 🧨 Diffusers implementa uma verificação de segurança: [`safety_checker`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/safety_checker.py) para prevenir conteúdo ofensivo ou nocivo, mas as capacidades de geração de imagem aprimorada do modelo podem ainda produzir conteúdo potencialmente nocivo. Para carregar o modelo com o método [`~DiffusionPipeline.from_pretrained`]: @@ -204,11 +198,8 @@ Para geração de exemplos reais, você precisará de um agendador para guiar o Agendadores gerenciam a retirada do ruído de uma amostra ruidosa para uma amostra menos ruidosa dado a saída do modelo - nesse caso, é o `noisy_residual`. - - -🧨 Diffusers é uma caixa de ferramentas para construir sistemas de difusão. Enquanto o [`DiffusionPipeline`] é uma forma conveniente de começar com um sistema de difusão pré-construído, você também pode escolher seus próprios modelos e agendadores separadamente para construir um sistema de difusão personalizado. - - +> [!TIP] +> 🧨 Diffusers é uma caixa de ferramentas para construir sistemas de difusão. Enquanto o [`DiffusionPipeline`] é uma forma conveniente de começar com um sistema de difusão pré-construído, você também pode escolher seus próprios modelos e agendadores separadamente para construir um sistema de difusão personalizado. Para o tour rápido, você irá instanciar o [`DDPMScheduler`] com o método [`~diffusers.ConfigMixin.from_config`]: @@ -232,11 +223,8 @@ DDPMScheduler { } ``` - - -💡 Perceba como o agendador é instanciado de uma configuração. Diferentemente de um modelo, um agendador não tem pesos treináveis e é livre de parâmetros! - - +> [!TIP] +> 💡 Perceba como o agendador é instanciado de uma configuração. Diferentemente de um modelo, um agendador não tem pesos treináveis e é livre de parâmetros! Um dos parâmetros mais importante são: diff --git a/docs/source/zh/conceptual/evaluation.md b/docs/source/zh/conceptual/evaluation.md index e809c8730d..770d197be0 100644 --- a/docs/source/zh/conceptual/evaluation.md +++ b/docs/source/zh/conceptual/evaluation.md @@ -92,11 +92,8 @@ images = sd_pipeline(sample_prompts, num_images_per_prompt=1, generator=generato 当使用多个待评估模型为所有提示词生成若干图像后，这些结果将提交给人类评估员进行打分。有关DrawBench和PartiPrompts基准测试的更多细节，请参阅各自的论文。 - - -在模型训练过程中查看推理样本有助于评估训练进度。我们的[训练脚本](https://github.com/huggingface/diffusers/tree/main/examples/)支持此功能，并额外提供TensorBoard和Weights & Biases日志记录功能。 - - +> [!TIP] +> 在模型训练过程中查看推理样本有助于评估训练进度。我们的[训练脚本](https://github.com/huggingface/diffusers/tree/main/examples/)支持此功能，并额外提供TensorBoard和Weights & Biases日志记录功能。 ## 定量评估 @@ -189,11 +186,8 @@ print(f"v-1-5版本的CLIP分数: {sd_clip_score_1_5}") 结果表明[v1-5](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5)检查点性能优于前代。但需注意，我们用于计算CLIP分数的提示词数量较少。实际评估时应使用更多样化且数量更大的提示词集。 - - -该分数存在固有局限性：训练数据中的标题是从网络爬取，并提取自图片关联的`alt`等标签。这些描述未必符合人类描述图像的方式，因此我们需要人工"设计"部分提示词。 - - +> [!WARNING] +> 该分数存在固有局限性：训练数据中的标题是从网络爬取，并提取自图片关联的`alt`等标签。这些描述未必符合人类描述图像的方式，因此我们需要人工"设计"部分提示词。 ### 图像条件式文本生成图像 @@ -402,11 +396,8 @@ print(f"CLIP方向相似度: {np.mean(scores)}") 该度量方法同样适用于类似流程，例如[`StableDiffusionPix2PixZeroPipeline`](https://huggingface.co/docs/diffusers/main/en/api/pipelines/pix2pix_zero#diffusers.StableDiffusionPix2PixZeroPipeline)。 - - -CLIP分数和CLIP方向相似度都依赖CLIP模型，可能导致评估结果存在偏差。 - - +> [!TIP] +> CLIP分数和CLIP方向相似度都依赖CLIP模型，可能导致评估结果存在偏差。 ***扩展IS、FID（后文讨论）或KID等指标存在困难***，当被评估模型是在大型图文数据集（如[LAION-5B数据集](https://laion.ai/blog/laion-5b/)）上预训练时。因为这些指标的底层都使用了在ImageNet-1k数据集上预训练的InceptionNet来提取图像特征。Stable Diffusion的预训练数据集与InceptionNet的预训练数据集可能重叠有限，因此不适合作为特征提取器。 @@ -536,19 +527,16 @@ FID分数越低越好。以下因素会影响FID结果： 对于最后两点，最佳实践是使用不同的随机种子和推理步数进行多次评估，然后报告平均结果。 - - -FID结果往往具有脆弱性，因为它依赖于许多因素： - -* 计算过程中使用的特定Inception模型 -* 计算实现的准确性 -* 图像格式（PNG和JPG的起点不同） - -需要注意的是，FID通常在比较相似实验时最有用，但除非作者仔细公开FID测量代码，否则很难复现论文结果。 - -这些注意事项同样适用于其他相关指标，如KID和IS。 - - +> [!WARNING] +> FID结果往往具有脆弱性，因为它依赖于许多因素： +> +> * 计算过程中使用的特定Inception模型 +> * 计算实现的准确性 +> * 图像格式（PNG和JPG的起点不同） +> +> 需要注意的是，FID通常在比较相似实验时最有用，但除非作者仔细公开FID测量代码，否则很难复现论文结果。 +> +> 这些注意事项同样适用于其他相关指标，如KID和IS。 最后，让我们可视化检查这些`fake_images`。 diff --git a/docs/source/zh/installation.md b/docs/source/zh/installation.md index fc77ea8c48..9941ed24ae 100644 --- a/docs/source/zh/installation.md +++ b/docs/source/zh/installation.md @@ -109,11 +109,8 @@ pip install -e ".[flax]" 现在，不只是在通常的库路径，Python 还会在你克隆的文件夹内寻找包。 例如，如果你的 Python 包通常安装在 `~/anaconda3/envs/main/lib/python3.10/Site-packages/`，Python 也会搜索你克隆到的文件夹。`~/diffusers/`。 - - -如果你想继续使用这个库，你必须保留 `diffusers` 文件夹。 - - +> [!WARNING] +> 如果你想继续使用这个库，你必须保留 `diffusers` 文件夹。 现在你可以用下面的命令轻松地将你克隆的 🤗 Diffusers 库更新到最新版本。 diff --git a/docs/source/zh/optimization/coreml.md b/docs/source/zh/optimization/coreml.md index 1d78866720..3926a5ddb0 100644 --- a/docs/source/zh/optimization/coreml.md +++ b/docs/source/zh/optimization/coreml.md @@ -13,11 +13,8 @@ http://www.apache.org/licenses/LICENSE-2.0 Core ML 模型可以利用 Apple 设备中所有可用的计算引擎：CPU、GPU 和 Apple Neural Engine（或 ANE，一种在 Apple Silicon Mac 和现代 iPhone/iPad 中可用的张量优化加速器）。根据模型及其运行的设备，Core ML 还可以混合和匹配计算引擎，例如，模型的某些部分可能在 CPU 上运行，而其他部分在 GPU 上运行。 - - -您还可以使用 PyTorch 内置的 `mps` 加速器在 Apple Silicon Mac 上运行 `diffusers` Python 代码库。这种方法在 [mps 指南](mps) 中有详细解释，但它与原生应用不兼容。 - - +> [!TIP] +> 您还可以使用 PyTorch 内置的 `mps` 加速器在 Apple Silicon Mac 上运行 `diffusers` Python 代码库。这种方法在 [mps 指南](mps) 中有详细解释，但它与原生应用不兼容。 ## Stable Diffusion Core ML 检查点 diff --git a/docs/source/zh/optimization/fp16.md b/docs/source/zh/optimization/fp16.md index 1088482d24..e1c4c7e57a 100644 --- a/docs/source/zh/optimization/fp16.md +++ b/docs/source/zh/optimization/fp16.md @@ -238,11 +238,8 @@ pipeline.unet = compile_regions(pipeline.unet, mode="reduce-overhead", fullgraph 一般来说，`sigmas`应该[保持在CPU上](https://github.com/huggingface/diffusers/blob/35a969d297cba69110d175ee79c59312b9f49e1e/src/diffusers/schedulers/scheduling_euler_discrete.py#L240)，以避免通信同步和延迟。 - - -参阅[torch.compile和Diffusers：峰值性能实践指南](https://pytorch.org/blog/torch-compile-and-diffusers-a-hands-on-guide-to-peak-performance/)博客文章，了解如何为扩散模型最大化`torch.compile`的性能。 - - +> [!TIP] +> 参阅[torch.compile和Diffusers：峰值性能实践指南](https://pytorch.org/blog/torch-compile-and-diffusers-a-hands-on-guide-to-peak-performance/)博客文章，了解如何为扩散模型最大化`torch.compile`的性能。 ### 基准测试 diff --git a/docs/source/zh/optimization/mps.md b/docs/source/zh/optimization/mps.md index c76a475336..48b08c5a12 100644 --- a/docs/source/zh/optimization/mps.md +++ b/docs/source/zh/optimization/mps.md @@ -35,11 +35,8 @@ image = pipe(prompt).images[0] image ``` - - -PyTorch [mps](https://pytorch.org/docs/stable/notes/mps.html) 后端不支持大小超过 `2**32` 的 NDArray。如果您遇到此问题，请提交 [Issue](https://github.com/huggingface/diffusers/issues/new/choose) 以便我们调查。 - - +> [!WARNING] +> PyTorch [mps](https://pytorch.org/docs/stable/notes/mps.html) 后端不支持大小超过 `2**32` 的 NDArray。如果您遇到此问题，请提交 [Issue](https://github.com/huggingface/diffusers/issues/new/choose) 以便我们调查。 如果您使用 **PyTorch 1.13**，您需要通过管道进行一次额外的"预热"传递。这是一个临时解决方法，用于解决首次推理传递产生的结果与后续传递略有不同的问题。您只需要执行此传递一次，并且在仅进行一次推理步骤后可以丢弃结果。 diff --git a/docs/source/zh/optimization/neuron.md b/docs/source/zh/optimization/neuron.md index 709404d56b..99d807a88c 100644 --- a/docs/source/zh/optimization/neuron.md +++ b/docs/source/zh/optimization/neuron.md @@ -17,11 +17,8 @@ Diffusers 功能可在 [AWS Inf2 实例](https://aws.amazon.com/ec2/instance-typ python -m pip install --upgrade-strategy eager optimum[neuronx] ``` - - -我们提供预构建的 [Hugging Face Neuron 深度学习 AMI](https://aws.amazon.com/marketplace/pp/prodview-gr3e6yiscria2)（DLAMI）和用于 Amazon SageMaker 的 Optimum Neuron 容器。建议正确设置您的环境。 - - +> [!TIP] +> 我们提供预构建的 [Hugging Face Neuron 深度学习 AMI](https://aws.amazon.com/marketplace/pp/prodview-gr3e6yiscria2)（DLAMI）和用于 Amazon SageMaker 的 Optimum Neuron 容器。建议正确设置您的环境。 下面的示例演示了如何在 inf2.8xlarge 实例上使用 Stable Diffusion XL 模型生成图像（一旦模型编译完成，您可以切换到更便宜的 inf2.xlarge 实例）。要生成一些图像，请使用 [`~optimum.neuron.NeuronStableDiffusionXLPipeline`] 类，该类类似于 Diffusers 中的 [`StableDiffusionXLPipeline`] 类。 diff --git a/docs/source/zh/optimization/onnx.md b/docs/source/zh/optimization/onnx.md index 4b3804d015..b70510d51b 100644 --- a/docs/source/zh/optimization/onnx.md +++ b/docs/source/zh/optimization/onnx.md @@ -31,11 +31,8 @@ image = pipeline(prompt).images[0] pipeline.save_pretrained("./onnx-stable-diffusion-v1-5") ``` - - -当前批量生成多个提示可能会占用过高内存。在问题修复前，建议采用迭代方式而非批量处理。 - - +> [!WARNING] +> 当前批量生成多个提示可能会占用过高内存。在问题修复前，建议采用迭代方式而非批量处理。 如需离线导出 ONNX 格式流水线供后续推理使用，请使用 [`optimum-cli export`](https://huggingface.co/docs/optimum/main/en/exporters/onnx/usage_guides/export_a_model#exporting-a-model-to-onnx-using-the-cli) 命令： diff --git a/docs/source/zh/optimization/xformers.md b/docs/source/zh/optimization/xformers.md index 9902feeee6..2a3a3d8341 100644 --- a/docs/source/zh/optimization/xformers.md +++ b/docs/source/zh/optimization/xformers.md @@ -17,16 +17,10 @@ http://www.apache.org/licenses/LICENSE-2.0 pip install xformers ``` - - -xFormers的`pip`安装包需要最新版本的PyTorch。如需使用旧版PyTorch，建议[从源码安装xFormers](https://github.com/facebookresearch/xformers#installing-xformers)。 - - +> [!TIP] +> xFormers的`pip`安装包需要最新版本的PyTorch。如需使用旧版PyTorch，建议[从源码安装xFormers](https://github.com/facebookresearch/xformers#installing-xformers)。 安装完成后，您可调用`enable_xformers_memory_efficient_attention()`来实现更快的推理速度和更低的内存占用，具体用法参见[此章节](memory#memory-efficient-attention)。 - - -根据[此问题](https://github.com/huggingface/diffusers/issues/2234#issuecomment-1416931212)反馈，xFormers `v0.0.16`版本在某些GPU上无法用于训练（微调或DreamBooth）。如遇此问题，请按照该issue评论区指引安装开发版本。 - - \ No newline at end of file +> [!WARNING] +> 根据[此问题](https://github.com/huggingface/diffusers/issues/2234#issuecomment-1416931212)反馈，xFormers `v0.0.16`版本在某些GPU上无法用于训练（微调或DreamBooth）。如遇此问题，请按照该issue评论区指引安装开发版本。 \ No newline at end of file diff --git a/docs/source/zh/quicktour.md b/docs/source/zh/quicktour.md index 08efaa87d2..2b8803384f 100644 --- a/docs/source/zh/quicktour.md +++ b/docs/source/zh/quicktour.md @@ -31,11 +31,8 @@ specific language governing permissions and limitations under the License. 快速入门将告诉你如何使用[`DiffusionPipeline`]进行推理，然后指导你如何结合模型和调度器以复现[`DiffusionPipeline`]内部发生的事情。 - - -快速入门是🧨[Diffusers入门](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/diffusers_intro.ipynb)的简化版，可以帮助你快速上手。如果你想了解更多关于🧨 Diffusers的目标、设计理念以及关于它的核心API的更多细节，可以点击🧨[Diffusers入门](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/diffusers_intro.ipynb)查看。 - - +> [!TIP] +> 快速入门是🧨[Diffusers入门](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/diffusers_intro.ipynb)的简化版，可以帮助你快速上手。如果你想了解更多关于🧨 Diffusers的目标、设计理念以及关于它的核心API的更多细节，可以点击🧨[Diffusers入门](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/diffusers_intro.ipynb)查看。 在开始之前，确认一下你已经安装好了所需要的库： @@ -66,11 +63,10 @@ pip install --upgrade diffusers accelerate transformers 您可以在Hugging Face Hub上使用[DiffusionPipeline]的任何检查点。 在本快速入门中，您将加载stable-diffusion-v1-5检查点，用于文本到图像生成。 -。 - -对于[Stable Diffusion](https://huggingface.co/CompVis/stable-diffusion)模型，在运行该模型之前，请先仔细阅读[许可证](https://huggingface.co/spaces/CompVis/stable-diffusion-license)。🧨 Diffusers实现了一个[`safety_checker`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/safety_checker.py)，以防止有攻击性的或有害的内容，但Stable Diffusion模型改进图像的生成能力仍有可能产生潜在的有害内容。 - - +> [!WARNING] +> 。 +> +> 对于[Stable Diffusion](https://huggingface.co/CompVis/stable-diffusion)模型，在运行该模型之前，请先仔细阅读[许可证](https://huggingface.co/spaces/CompVis/stable-diffusion-license)。🧨 Diffusers实现了一个[`safety_checker`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/safety_checker.py)，以防止有攻击性的或有害的内容，但Stable Diffusion模型改进图像的生成能力仍有可能产生潜在的有害内容。 用[`~DiffusionPipeline.from_pretrained`]方法加载模型。 @@ -221,11 +217,8 @@ torch.Size([1, 3, 256, 256]) - - -🧨 Diffusers是一个用于构建扩散系统的工具箱。预定义好的扩散系统[`DiffusionPipeline`]能方便你快速试用，你也可以单独选择自己的模型和调度器组件来建立一个自定义的扩散系统。 - - +> [!TIP] +> 🧨 Diffusers是一个用于构建扩散系统的工具箱。预定义好的扩散系统[`DiffusionPipeline`]能方便你快速试用，你也可以单独选择自己的模型和调度器组件来建立一个自定义的扩散系统。 在快速入门教程中，你将用它的[`~diffusers.ConfigMixin.from_config`]方法实例化[`DDPMScheduler`]： @@ -249,12 +242,8 @@ DDPMScheduler { } ``` - - - -💡 注意调度器是如何从配置中实例化的。与模型不同，调度器没有可训练的权重，而且是无参数的。 - - +> [!TIP] +> 💡 注意调度器是如何从配置中实例化的。与模型不同，调度器没有可训练的权重，而且是无参数的。 * `num_train_timesteps`：去噪过程的长度，或者换句话说，将随机高斯噪声处理成数据样本所需的时间步数。 * `beta_schedule`：用于推理和训练的噪声表。 diff --git a/docs/source/zh/stable_diffusion.md b/docs/source/zh/stable_diffusion.md index bf9288c5b7..d337fb41a0 100644 --- a/docs/source/zh/stable_diffusion.md +++ b/docs/source/zh/stable_diffusion.md @@ -1,264 +1,258 @@ - - -# 有效且高效的扩散 - -[[open-in-colab]] - -让 [`DiffusionPipeline`] 生成特定风格或包含你所想要的内容的图像可能会有些棘手。 通常情况下，你需要多次运行 [`DiffusionPipeline`] 才能得到满意的图像。但是从无到有生成图像是一个计算密集的过程，特别是如果你要一遍又一遍地进行推理运算。 - -这就是为什么从pipeline中获得最高的 *computational* (speed) 和 *memory* (GPU RAM) 非常重要 ，以减少推理周期之间的时间，从而使迭代速度更快。 - - -本教程将指导您如何通过 [`DiffusionPipeline`] 更快、更好地生成图像。 - - -首先，加载 [`stable-diffusion-v1-5/stable-diffusion-v1-5`](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) 模型: - -```python -from diffusers import DiffusionPipeline - -model_id = "stable-diffusion-v1-5/stable-diffusion-v1-5" -pipeline = DiffusionPipeline.from_pretrained(model_id, use_safetensors=True) -``` - -本教程将使用的提示词是 [`portrait photo of a old warrior chief`] ，但是你可以随心所欲的想象和构造自己的提示词： - -```python -prompt = "portrait photo of a old warrior chief" -``` - -## 速度 - - - -💡 如果你没有 GPU, 你可以从像 [Colab](https://colab.research.google.com/) 这样的 GPU 提供商获取免费的 GPU ! - - - -加速推理的最简单方法之一是将 pipeline 放在 GPU 上 ，就像使用任何 PyTorch 模块一样： - -```python -pipeline = pipeline.to("cuda") -``` - -为了确保您可以使用相同的图像并对其进行改进，使用 [`Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) 方法，然后设置一个随机数种子 以确保其 [复现性](./using-diffusers/reusing_seeds): - -```python -import torch - -generator = torch.Generator("cuda").manual_seed(0) -``` - -现在，你可以生成一个图像： - -```python -image = pipeline(prompt, generator=generator).images[0] -image -``` - -
- -
- -在 T4 GPU 上，这个过程大概要30秒（如果你的 GPU 比 T4 好，可能会更快）。在默认情况下，[`DiffusionPipeline`] 使用完整的 `float32` 精度进行 50 步推理。你可以通过降低精度（如 `float16` ）或者减少推理步数来加速整个过程 - - -让我们把模型的精度降低至 `float16` ，然后生成一张图像： - -```python -import torch - -pipeline = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, use_safetensors=True) -pipeline = pipeline.to("cuda") -generator = torch.Generator("cuda").manual_seed(0) -image = pipeline(prompt, generator=generator).images[0] -image -``` - -
- -
- -这一次，生成图像只花了约 11 秒，比之前快了近 3 倍！ - - - -💡 我们强烈建议把 pipeline 精度降低至 `float16` , 到目前为止, 我们很少看到输出质量有任何下降。 - - - -另一个选择是减少推理步数。 你可以选择一个更高效的调度器 (*scheduler*) 可以减少推理步数同时保证输出质量。您可以在 [DiffusionPipeline] 中通过调用compatibles方法找到与当前模型兼容的调度器 (*scheduler*)。 - -```python -pipeline.scheduler.compatibles -[ - diffusers.schedulers.scheduling_lms_discrete.LMSDiscreteScheduler, - diffusers.schedulers.scheduling_unipc_multistep.UniPCMultistepScheduler, - diffusers.schedulers.scheduling_k_dpm_2_discrete.KDPM2DiscreteScheduler, - diffusers.schedulers.scheduling_deis_multistep.DEISMultistepScheduler, - diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler, - diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler, - diffusers.schedulers.scheduling_ddpm.DDPMScheduler, - diffusers.schedulers.scheduling_dpmsolver_singlestep.DPMSolverSinglestepScheduler, - diffusers.schedulers.scheduling_k_dpm_2_ancestral_discrete.KDPM2AncestralDiscreteScheduler, - diffusers.schedulers.scheduling_heun_discrete.HeunDiscreteScheduler, - diffusers.schedulers.scheduling_pndm.PNDMScheduler, - diffusers.schedulers.scheduling_euler_ancestral_discrete.EulerAncestralDiscreteScheduler, - diffusers.schedulers.scheduling_ddim.DDIMScheduler, -] -``` - -Stable Diffusion 模型默认使用的是 [`PNDMScheduler`] ，通常要大概50步推理, 但是像 [`DPMSolverMultistepScheduler`] 这样更高效的调度器只要大概 20 或 25 步推理. 使用 [`ConfigMixin.from_config`] 方法加载新的调度器: - -```python -from diffusers import DPMSolverMultistepScheduler - -pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config) -``` - -现在将 `num_inference_steps` 设置为 20: - -```python -generator = torch.Generator("cuda").manual_seed(0) -image = pipeline(prompt, generator=generator, num_inference_steps=20).images[0] -image -``` - -
- -
- -太棒了！你成功把推理时间缩短到 4 秒！⚡️ - -## 内存 - -改善 pipeline 性能的另一个关键是减少内存的使用量，这间接意味着速度更快，因为你经常试图最大化每秒生成的图像数量。要想知道你一次可以生成多少张图片，最简单的方法是尝试不同的batch size，直到出现`OutOfMemoryError` (OOM)。 - -创建一个函数，为每一批要生成的图像分配提示词和 `Generators` 。请务必为每个`Generator` 分配一个种子，以便于复现良好的结果。 - - -```python -def get_inputs(batch_size=1): - generator = [torch.Generator("cuda").manual_seed(i) for i in range(batch_size)] - prompts = batch_size * [prompt] - num_inference_steps = 20 - - return {"prompt": prompts, "generator": generator, "num_inference_steps": num_inference_steps} -``` - -设置 `batch_size=4` ，然后看一看我们消耗了多少内存: - -```python -from diffusers.utils import make_image_grid - -images = pipeline(**get_inputs(batch_size=4)).images -make_image_grid(images, 2, 2) -``` - -除非你有一个更大内存的GPU, 否则上述代码会返回 `OOM` 错误! 大部分内存被 cross-attention 层使用。按顺序运行可以节省大量内存，而不是在批处理中进行。你可以为 pipeline 配置 [`~DiffusionPipeline.enable_attention_slicing`] 函数: - -```python -pipeline.enable_attention_slicing() -``` - -现在尝试把 `batch_size` 增加到 8! - -```python -images = pipeline(**get_inputs(batch_size=8)).images -make_image_grid(images, rows=2, cols=4) -``` - -
- -
- -以前你不能一批生成 4 张图片，而现在你可以在一张图片里面生成八张图片而只需要大概3.5秒！这可能是 T4 GPU 在不牺牲质量的情况运行速度最快的一种方法。 - -## 质量 - -在最后两节中, 你要学习如何通过 `fp16` 来优化 pipeline 的速度, 通过使用性能更高的调度器来减少推理步数, 使用注意力切片（*enabling attention slicing*）方法来节省内存。现在，你将关注的是如何提高图像的质量。 - -### 更好的 checkpoints - -有个显而易见的方法是使用更好的 checkpoints。 Stable Diffusion 模型是一个很好的起点, 自正式发布以来，还发布了几个改进版本。然而, 使用更新的版本并不意味着你会得到更好的结果。你仍然需要尝试不同的 checkpoints ，并做一些研究 (例如使用 [negative prompts](https://minimaxir.com/2022/11/stable-diffusion-negative-prompt/)) 来获得更好的结果。 - -随着该领域的发展, 有越来越多经过微调的高质量的 checkpoints 用来生成不一样的风格. 在 [Hub](https://huggingface.co/models?library=diffusers&sort=downloads) 和 [Diffusers Gallery](https://huggingface.co/spaces/huggingface-projects/diffusers-gallery) 寻找你感兴趣的一种! - -### 更好的 pipeline 组件 - -也可以尝试用新版本替换当前 pipeline 组件。让我们加载最新的 [autodecoder](https://huggingface.co/stabilityai/stable-diffusion-2-1/tree/main/vae) 从 Stability AI 加载到 pipeline, 并生成一些图像: - -```python -from diffusers import AutoencoderKL - -vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16).to("cuda") -pipeline.vae = vae -images = pipeline(**get_inputs(batch_size=8)).images -make_image_grid(images, rows=2, cols=4) -``` - -
- -
- -### 更好的提示词工程 - -用于生成图像的文本非常重要, 因此被称为 *提示词工程*。 在设计提示词工程应注意如下事项: - -- 我想生成的图像或类似图像如何存储在互联网上？ -- 我可以提供哪些额外的细节来引导模型朝着我想要的风格生成？ - -考虑到这一点，让我们改进提示词，以包含颜色和更高质量的细节： - -```python -prompt += ", tribal panther make up, blue on red, side profile, looking away, serious eyes" -prompt += " 50mm portrait photography, hard rim lighting photography--beta --ar 2:3 --beta --upbeta" -``` - -使用新的提示词生成一批图像: - -```python -images = pipeline(**get_inputs(batch_size=8)).images -make_image_grid(images, rows=2, cols=4) -``` - -
- -
- -非常的令人印象深刻! Let's tweak the second image - 把 `Generator` 的种子设置为 `1` - 添加一些关于年龄的主题文本: - -```python -prompts = [ - "portrait photo of the oldest warrior chief, tribal panther make up, blue on red, side profile, looking away, serious eyes 50mm portrait photography, hard rim lighting photography--beta --ar 2:3 --beta --upbeta", - "portrait photo of a old warrior chief, tribal panther make up, blue on red, side profile, looking away, serious eyes 50mm portrait photography, hard rim lighting photography--beta --ar 2:3 --beta --upbeta", - "portrait photo of a warrior chief, tribal panther make up, blue on red, side profile, looking away, serious eyes 50mm portrait photography, hard rim lighting photography--beta --ar 2:3 --beta --upbeta", - "portrait photo of a young warrior chief, tribal panther make up, blue on red, side profile, looking away, serious eyes 50mm portrait photography, hard rim lighting photography--beta --ar 2:3 --beta --upbeta", -] - -generator = [torch.Generator("cuda").manual_seed(1) for _ in range(len(prompts))] -images = pipeline(prompt=prompts, generator=generator, num_inference_steps=25).images -make_image_grid(images, 2, 2) -``` - -
- -
- -## 最后 - -在本教程中, 您学习了如何优化[`DiffusionPipeline`]以提高计算和内存效率，以及提高生成输出的质量. 如果你有兴趣让你的 pipeline 更快, 可以看一看以下资源: - -- 学习 [PyTorch 2.0](./optimization/torch2.0) 和 [`torch.compile`](https://pytorch.org/docs/stable/generated/torch.compile.html) 可以让推理速度提高 5 - 300% . 在 A100 GPU 上, 推理速度可以提高 50% ! -- 如果你没法用 PyTorch 2, 我们建议你安装 [xFormers](./optimization/xformers)。它的内存高效注意力机制（*memory-efficient attention mechanism*）与PyTorch 1.13.1配合使用，速度更快，内存消耗更少。 -- 其他的优化技术, 如：模型卸载（*model offloading*）, 包含在 [这份指南](./optimization/fp16). + + +# 有效且高效的扩散 + +[[open-in-colab]] + +让 [`DiffusionPipeline`] 生成特定风格或包含你所想要的内容的图像可能会有些棘手。 通常情况下，你需要多次运行 [`DiffusionPipeline`] 才能得到满意的图像。但是从无到有生成图像是一个计算密集的过程，特别是如果你要一遍又一遍地进行推理运算。 + +这就是为什么从pipeline中获得最高的 *computational* (speed) 和 *memory* (GPU RAM) 非常重要 ，以减少推理周期之间的时间，从而使迭代速度更快。 + + +本教程将指导您如何通过 [`DiffusionPipeline`] 更快、更好地生成图像。 + + +首先，加载 [`stable-diffusion-v1-5/stable-diffusion-v1-5`](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) 模型: + +```python +from diffusers import DiffusionPipeline + +model_id = "stable-diffusion-v1-5/stable-diffusion-v1-5" +pipeline = DiffusionPipeline.from_pretrained(model_id, use_safetensors=True) +``` + +本教程将使用的提示词是 [`portrait photo of a old warrior chief`] ，但是你可以随心所欲的想象和构造自己的提示词： + +```python +prompt = "portrait photo of a old warrior chief" +``` + +## 速度 + +> [!TIP] +> 💡 如果你没有 GPU, 你可以从像 [Colab](https://colab.research.google.com/) 这样的 GPU 提供商获取免费的 GPU ! + +加速推理的最简单方法之一是将 pipeline 放在 GPU 上 ，就像使用任何 PyTorch 模块一样： + +```python +pipeline = pipeline.to("cuda") +``` + +为了确保您可以使用相同的图像并对其进行改进，使用 [`Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) 方法，然后设置一个随机数种子 以确保其 [复现性](./using-diffusers/reusing_seeds): + +```python +import torch + +generator = torch.Generator("cuda").manual_seed(0) +``` + +现在，你可以生成一个图像： + +```python +image = pipeline(prompt, generator=generator).images[0] +image +``` + +
+ +
+ +在 T4 GPU 上，这个过程大概要30秒（如果你的 GPU 比 T4 好，可能会更快）。在默认情况下，[`DiffusionPipeline`] 使用完整的 `float32` 精度进行 50 步推理。你可以通过降低精度（如 `float16` ）或者减少推理步数来加速整个过程 + + +让我们把模型的精度降低至 `float16` ，然后生成一张图像： + +```python +import torch + +pipeline = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, use_safetensors=True) +pipeline = pipeline.to("cuda") +generator = torch.Generator("cuda").manual_seed(0) +image = pipeline(prompt, generator=generator).images[0] +image +``` + +
+ +
+ +这一次，生成图像只花了约 11 秒，比之前快了近 3 倍！ + +> [!TIP] +> 💡 我们强烈建议把 pipeline 精度降低至 `float16` , 到目前为止, 我们很少看到输出质量有任何下降。 + +另一个选择是减少推理步数。 你可以选择一个更高效的调度器 (*scheduler*) 可以减少推理步数同时保证输出质量。您可以在 [DiffusionPipeline] 中通过调用compatibles方法找到与当前模型兼容的调度器 (*scheduler*)。 + +```python +pipeline.scheduler.compatibles +[ + diffusers.schedulers.scheduling_lms_discrete.LMSDiscreteScheduler, + diffusers.schedulers.scheduling_unipc_multistep.UniPCMultistepScheduler, + diffusers.schedulers.scheduling_k_dpm_2_discrete.KDPM2DiscreteScheduler, + diffusers.schedulers.scheduling_deis_multistep.DEISMultistepScheduler, + diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler, + diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler, + diffusers.schedulers.scheduling_ddpm.DDPMScheduler, + diffusers.schedulers.scheduling_dpmsolver_singlestep.DPMSolverSinglestepScheduler, + diffusers.schedulers.scheduling_k_dpm_2_ancestral_discrete.KDPM2AncestralDiscreteScheduler, + diffusers.schedulers.scheduling_heun_discrete.HeunDiscreteScheduler, + diffusers.schedulers.scheduling_pndm.PNDMScheduler, + diffusers.schedulers.scheduling_euler_ancestral_discrete.EulerAncestralDiscreteScheduler, + diffusers.schedulers.scheduling_ddim.DDIMScheduler, +] +``` + +Stable Diffusion 模型默认使用的是 [`PNDMScheduler`] ，通常要大概50步推理, 但是像 [`DPMSolverMultistepScheduler`] 这样更高效的调度器只要大概 20 或 25 步推理. 使用 [`ConfigMixin.from_config`] 方法加载新的调度器: + +```python +from diffusers import DPMSolverMultistepScheduler + +pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config) +``` + +现在将 `num_inference_steps` 设置为 20: + +```python +generator = torch.Generator("cuda").manual_seed(0) +image = pipeline(prompt, generator=generator, num_inference_steps=20).images[0] +image +``` + +
+ +
+ +太棒了！你成功把推理时间缩短到 4 秒！⚡️ + +## 内存 + +改善 pipeline 性能的另一个关键是减少内存的使用量，这间接意味着速度更快，因为你经常试图最大化每秒生成的图像数量。要想知道你一次可以生成多少张图片，最简单的方法是尝试不同的batch size，直到出现`OutOfMemoryError` (OOM)。 + +创建一个函数，为每一批要生成的图像分配提示词和 `Generators` 。请务必为每个`Generator` 分配一个种子，以便于复现良好的结果。 + + +```python +def get_inputs(batch_size=1): + generator = [torch.Generator("cuda").manual_seed(i) for i in range(batch_size)] + prompts = batch_size * [prompt] + num_inference_steps = 20 + + return {"prompt": prompts, "generator": generator, "num_inference_steps": num_inference_steps} +``` + +设置 `batch_size=4` ，然后看一看我们消耗了多少内存: + +```python +from diffusers.utils import make_image_grid + +images = pipeline(**get_inputs(batch_size=4)).images +make_image_grid(images, 2, 2) +``` + +除非你有一个更大内存的GPU, 否则上述代码会返回 `OOM` 错误! 大部分内存被 cross-attention 层使用。按顺序运行可以节省大量内存，而不是在批处理中进行。你可以为 pipeline 配置 [`~DiffusionPipeline.enable_attention_slicing`] 函数: + +```python +pipeline.enable_attention_slicing() +``` + +现在尝试把 `batch_size` 增加到 8! + +```python +images = pipeline(**get_inputs(batch_size=8)).images +make_image_grid(images, rows=2, cols=4) +``` + +
+ +
+ +以前你不能一批生成 4 张图片，而现在你可以在一张图片里面生成八张图片而只需要大概3.5秒！这可能是 T4 GPU 在不牺牲质量的情况运行速度最快的一种方法。 + +## 质量 + +在最后两节中, 你要学习如何通过 `fp16` 来优化 pipeline 的速度, 通过使用性能更高的调度器来减少推理步数, 使用注意力切片（*enabling attention slicing*）方法来节省内存。现在，你将关注的是如何提高图像的质量。 + +### 更好的 checkpoints + +有个显而易见的方法是使用更好的 checkpoints。 Stable Diffusion 模型是一个很好的起点, 自正式发布以来，还发布了几个改进版本。然而, 使用更新的版本并不意味着你会得到更好的结果。你仍然需要尝试不同的 checkpoints ，并做一些研究 (例如使用 [negative prompts](https://minimaxir.com/2022/11/stable-diffusion-negative-prompt/)) 来获得更好的结果。 + +随着该领域的发展, 有越来越多经过微调的高质量的 checkpoints 用来生成不一样的风格. 在 [Hub](https://huggingface.co/models?library=diffusers&sort=downloads) 和 [Diffusers Gallery](https://huggingface.co/spaces/huggingface-projects/diffusers-gallery) 寻找你感兴趣的一种! + +### 更好的 pipeline 组件 + +也可以尝试用新版本替换当前 pipeline 组件。让我们加载最新的 [autodecoder](https://huggingface.co/stabilityai/stable-diffusion-2-1/tree/main/vae) 从 Stability AI 加载到 pipeline, 并生成一些图像: + +```python +from diffusers import AutoencoderKL + +vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16).to("cuda") +pipeline.vae = vae +images = pipeline(**get_inputs(batch_size=8)).images +make_image_grid(images, rows=2, cols=4) +``` + +
+ +
+ +### 更好的提示词工程 + +用于生成图像的文本非常重要, 因此被称为 *提示词工程*。 在设计提示词工程应注意如下事项: + +- 我想生成的图像或类似图像如何存储在互联网上？ +- 我可以提供哪些额外的细节来引导模型朝着我想要的风格生成？ + +考虑到这一点，让我们改进提示词，以包含颜色和更高质量的细节： + +```python +prompt += ", tribal panther make up, blue on red, side profile, looking away, serious eyes" +prompt += " 50mm portrait photography, hard rim lighting photography--beta --ar 2:3 --beta --upbeta" +``` + +使用新的提示词生成一批图像: + +```python +images = pipeline(**get_inputs(batch_size=8)).images +make_image_grid(images, rows=2, cols=4) +``` + +
+ +
+ +非常的令人印象深刻! Let's tweak the second image - 把 `Generator` 的种子设置为 `1` - 添加一些关于年龄的主题文本: + +```python +prompts = [ + "portrait photo of the oldest warrior chief, tribal panther make up, blue on red, side profile, looking away, serious eyes 50mm portrait photography, hard rim lighting photography--beta --ar 2:3 --beta --upbeta", + "portrait photo of a old warrior chief, tribal panther make up, blue on red, side profile, looking away, serious eyes 50mm portrait photography, hard rim lighting photography--beta --ar 2:3 --beta --upbeta", + "portrait photo of a warrior chief, tribal panther make up, blue on red, side profile, looking away, serious eyes 50mm portrait photography, hard rim lighting photography--beta --ar 2:3 --beta --upbeta", + "portrait photo of a young warrior chief, tribal panther make up, blue on red, side profile, looking away, serious eyes 50mm portrait photography, hard rim lighting photography--beta --ar 2:3 --beta --upbeta", +] + +generator = [torch.Generator("cuda").manual_seed(1) for _ in range(len(prompts))] +images = pipeline(prompt=prompts, generator=generator, num_inference_steps=25).images +make_image_grid(images, 2, 2) +``` + +
+ +
+ +## 最后 + +在本教程中, 您学习了如何优化[`DiffusionPipeline`]以提高计算和内存效率，以及提高生成输出的质量. 如果你有兴趣让你的 pipeline 更快, 可以看一看以下资源: + +- 学习 [PyTorch 2.0](./optimization/torch2.0) 和 [`torch.compile`](https://pytorch.org/docs/stable/generated/torch.compile.html) 可以让推理速度提高 5 - 300% . 在 A100 GPU 上, 推理速度可以提高 50% ! +- 如果你没法用 PyTorch 2, 我们建议你安装 [xFormers](./optimization/xformers)。它的内存高效注意力机制（*memory-efficient attention mechanism*）与PyTorch 1.13.1配合使用，速度更快，内存消耗更少。 +- 其他的优化技术, 如：模型卸载（*model offloading*）, 包含在 [这份指南](./optimization/fp16). diff --git a/docs/source/zh/training/controlnet.md b/docs/source/zh/training/controlnet.md index e943177ced..84bc3263a8 100644 --- a/docs/source/zh/training/controlnet.md +++ b/docs/source/zh/training/controlnet.md @@ -68,11 +68,8 @@ pip install -r requirements_flax.txt - - -🤗 Accelerate 是一个支持多GPU/TPU训练和混合精度的库，它能根据硬件环境自动配置训练方案。参阅 🤗 Accelerate [快速入门](https://huggingface.co/docs/accelerate/quicktour) 了解更多。 - - +> [!TIP] +> 🤗 Accelerate 是一个支持多GPU/TPU训练和混合精度的库，它能根据硬件环境自动配置训练方案。参阅 🤗 Accelerate [快速入门](https://huggingface.co/docs/accelerate/quicktour) 了解更多。 初始化🤗 Accelerate环境： @@ -96,11 +93,8 @@ write_basic_config() 最后，如需训练自定义数据集，请参阅 [创建训练数据集](create_dataset) 指南了解数据准备方法。 - - -下文重点解析脚本中的关键模块，但不会覆盖所有实现细节。如需深入了解，建议直接阅读 [脚本源码](https://github.com/huggingface/diffusers/blob/main/examples/controlnet/train_controlnet.py)，如有疑问欢迎反馈。 - - +> [!TIP] +> 下文重点解析脚本中的关键模块，但不会覆盖所有实现细节。如需深入了解，建议直接阅读 [脚本源码](https://github.com/huggingface/diffusers/blob/main/examples/controlnet/train_controlnet.py)，如有疑问欢迎反馈。 ## 脚本参数 @@ -135,11 +129,8 @@ accelerate launch train_controlnet.py \ 脚本中的 [`make_train_dataset`](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/controlnet/train_controlnet.py#L582) 函数负责数据预处理，除常规的文本标注分词和图像变换外，还包含条件图像的特效处理： - - -在TPU上流式加载数据集时，🤗 Datasets库可能成为性能瓶颈（因其未针对图像数据优化）。建议考虑 [WebDataset](https://webdataset.github.io/webdataset/)、[TorchData](https://github.com/pytorch/data) 或 [TensorFlow Datasets](https://www.tensorflow.org/datasets/tfless_tfds) 等高效数据格式。 - - +> [!TIP] +> 在TPU上流式加载数据集时，🤗 Datasets库可能成为性能瓶颈（因其未针对图像数据优化）。建议考虑 [WebDataset](https://webdataset.github.io/webdataset/)、[TorchData](https://github.com/pytorch/data) 或 [TensorFlow Datasets](https://www.tensorflow.org/datasets/tfless_tfds) 等高效数据格式。 ```py conditioning_image_transforms = transforms.Compose( @@ -304,11 +295,8 @@ tensorboard --logdir runs/fill-circle-100steps-20230411_165612/ 在 [http://localhost:6006/#profile](http://localhost:6006/#profile) 查看分析结果。 - - -若遇到插件版本冲突，建议重新安装TensorFlow和Tensorboard。注意性能分析插件仍处实验阶段，部分视图可能不完整。`trace_viewer` 会截断超过1M的事件记录，在编译步骤分析时可能导致设备轨迹丢失。 - - +> [!WARNING] +> 若遇到插件版本冲突，建议重新安装TensorFlow和Tensorboard。注意性能分析插件仍处实验阶段，部分视图可能不完整。`trace_viewer` 会截断超过1M的事件记录，在编译步骤分析时可能导致设备轨迹丢失。 ```bash python3 train_controlnet_flax.py \ diff --git a/docs/source/zh/training/distributed_inference.md b/docs/source/zh/training/distributed_inference.md index e0537735b2..60297371d6 100644 --- a/docs/source/zh/training/distributed_inference.md +++ b/docs/source/zh/training/distributed_inference.md @@ -43,11 +43,8 @@ with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt: accelerate launch run_distributed.py --num_processes=2 ``` - - -参考这个最小示例 [脚本](https://gist.github.com/sayakpaul/cfaebd221820d7b43fae638b4dfa01ba) 以在多个 GPU 上运行推理。要了解更多信息，请查看 [使用 🤗 Accelerate 进行分布式推理](https://huggingface.co/docs/accelerate/en/usage_guides/distributed_inference#distributed-inference-with-accelerate) 指南。 - - +> [!TIP] +> 参考这个最小示例 [脚本](https://gist.github.com/sayakpaul/cfaebd221820d7b43fae638b4dfa01ba) 以在多个 GPU 上运行推理。要了解更多信息，请查看 [使用 🤗 Accelerate 进行分布式推理](https://huggingface.co/docs/accelerate/en/usage_guides/distributed_inference#distributed-inference-with-accelerate) 指南。 ## PyTorch Distributed diff --git a/docs/source/zh/training/dreambooth.md b/docs/source/zh/training/dreambooth.md index 493c5385ff..cae5e30be0 100644 --- a/docs/source/zh/training/dreambooth.md +++ b/docs/source/zh/training/dreambooth.md @@ -44,11 +44,8 @@ pip install -r requirements_flax.txt - - -🤗 Accelerate 是一个库，用于帮助您在多个 GPU/TPU 上或使用混合精度进行训练。它会根据您的硬件和环境自动配置训练设置。查看 🤗 Accelerate [快速入门](https://huggingface.co/docs/accelerate/quicktour) 以了解更多信息。 - - +> [!TIP] +> 🤗 Accelerate 是一个库，用于帮助您在多个 GPU/TPU 上或使用混合精度进行训练。它会根据您的硬件和环境自动配置训练设置。查看 🤗 Accelerate [快速入门](https://huggingface.co/docs/accelerate/quicktour) 以了解更多信息。 初始化 🤗 Accelerate 环境： @@ -73,19 +70,13 @@ write_basic_config() 最后，如果您想在自己的数据集上训练模型，请查看 [创建用于训练的数据集](create_dataset) 指南，了解如何创建与 训练脚本。 - - -以下部分重点介绍了训练脚本中对于理解如何修改它很重要的部分，但并未详细涵盖脚本的每个方面。如果您有兴趣了解更多，请随时阅读[脚本](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth.py)，并告诉我们如果您有任何问题或疑虑。 - - +> [!TIP] +> 以下部分重点介绍了训练脚本中对于理解如何修改它很重要的部分，但并未详细涵盖脚本的每个方面。如果您有兴趣了解更多，请随时阅读[脚本](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth.py)，并告诉我们如果您有任何问题或疑虑。 ## 脚本参数 - - -DreamBooth 对训练超参数非常敏感，容易过拟合。阅读 [使用 🧨 Diffusers 训练 Stable Diffusion 与 Dreambooth](https://huggingface.co/blog/dreambooth) 博客文章，了解针对不同主题的推荐设置，以帮助您选择合适的超参数。 - - +> [!WARNING] +> DreamBooth 对训练超参数非常敏感，容易过拟合。阅读 [使用 🧨 Diffusers 训练 Stable Diffusion 与 Dreambooth](https://huggingface.co/blog/dreambooth) 博客文章，了解针对不同主题的推荐设置，以帮助您选择合适的超参数。 训练脚本提供了许多参数来自定义您的训练运行。所有参数及其描述都可以在 [`parse_args()`](https://github.com/huggingface/diffusers/blob/072e00897a7cf4302c347a63ec917b4b8add16d4/examples/dreambooth/train_dreambooth.py#L228) 函数中找到。参数设置了默认值，这些默认值应该开箱即用效果不错，但如果您愿意，也可以在训练命令中设置自己的值。 @@ -359,29 +350,26 @@ python train_dreambooth_flax.py \ 训练完成后，您可以使用新训练的模型进行推理！ - - -等不及在训练完成前就尝试您的模型进行推理？🤭 请确保安装了最新版本的 🤗 Accelerate。 - -```py -from diffusers import DiffusionPipeline, UNet2DConditionModel -from transformers import CLIPTextModel -import torch - -unet = UNet2DConditionModel.from_pretrained("path/to/model/checkpoint-100/unet") - -# 如果您使用了 `--args.train_text_encoder` 进行训练，请确保也加载文本编码器 -text_encoder = CLIPTextModel.from_pretrained("path/to/model/checkpoint-100/checkpoint-100/text_encoder") - -pipeline = DiffusionPipeline.from_pretrained( - "stable-diffusion-v1-5/stable-diffusion-v1-5", unet=unet, text_encoder=text_encoder, dtype=torch.float16, -).to("cuda") - -image = pipeline("A photo of sks dog in a bucket", num_inference_steps=50, guidance_scale=7.5).images[0] -image.save("dog-bucket.png") -``` - - +> [!TIP] +> 等不及在训练完成前就尝试您的模型进行推理？🤭 请确保安装了最新版本的 🤗 Accelerate。 +> +> ```py +> from diffusers import DiffusionPipeline, UNet2DConditionModel +> from transformers import CLIPTextModel +> import torch +> +> unet = UNet2DConditionModel.from_pretrained("path/to/model/checkpoint-100/unet") +> +> # 如果您使用了 `--args.train_text_encoder` 进行训练，请确保也加载文本编码器 +> text_encoder = CLIPTextModel.from_pretrained("path/to/model/checkpoint-100/checkpoint-100/text_encoder") +> +> pipeline = DiffusionPipeline.from_pretrained( +> "stable-diffusion-v1-5/stable-diffusion-v1-5", unet=unet, text_encoder=text_encoder, dtype=torch.float16, +> ).to("cuda") +> +> image = pipeline("A photo of sks dog in a bucket", num_inference_steps=50, guidance_scale=7.5).images[0] +> image.save("dog-bucket.png") +> ``` diff --git a/docs/source/zh/training/instructpix2pix.md b/docs/source/zh/training/instructpix2pix.md index b1b616366a..1f9f4eb21e 100644 --- a/docs/source/zh/training/instructpix2pix.md +++ b/docs/source/zh/training/instructpix2pix.md @@ -31,11 +31,8 @@ cd examples/instruct_pix2pix pip install -r requirements.txt ``` - - -🤗 Accelerate 是一个库，用于帮助您在多个 GPU/TPU 上或使用混合精度进行训练。它将根据您的硬件和环境自动配置训练设置。查看 🤗 Accelerate [快速导览](https://huggingface.co/docs/accelerate/quicktour) 以了解更多信息。 - - +> [!TIP] +> 🤗 Accelerate 是一个库，用于帮助您在多个 GPU/TPU 上或使用混合精度进行训练。它将根据您的硬件和环境自动配置训练设置。查看 🤗 Accelerate [快速导览](https://huggingface.co/docs/accelerate/quicktour) 以了解更多信息。 初始化一个 🤗 Accelerate 环境： @@ -59,11 +56,8 @@ write_basic_config() 最后，如果您想在自己的数据集上训练模型，请查看 [创建用于训练的数据集](create_dataset) 指南，了解如何创建与训练脚本兼容的数据集。 - - -以下部分重点介绍了训练脚本中对于理解如何修改它很重要的部分，但并未详细涵盖脚本的每个方面。如果您有兴趣了解更多，请随时阅读 [脚本](https://github.com/huggingface/diffusers/blob/main/examples/instruct_pix2pix/train_instruct_pix2pix.py)，并告诉我们如果您有任何问题或疑虑。 - - +> [!TIP] +> 以下部分重点介绍了训练脚本中对于理解如何修改它很重要的部分，但并未详细涵盖脚本的每个方面。如果您有兴趣了解更多，请随时阅读 [脚本](https://github.com/huggingface/diffusers/blob/main/examples/instruct_pix2pix/train_instruct_pix2pix.py)，并告诉我们如果您有任何问题或疑虑。 ## 脚本参数 @@ -176,15 +170,12 @@ if args.conditioning_dropout_prob is not None: 将 `MODEL_NAME` 环境变量设置为模型名称（可以是 Hub 上的模型 ID 或本地模型的路径），并将 `DATASET_ID` 设置为 Hub 上数据集的名称。脚本会创建并保存所有组件（特征提取器、调度器、文本编码器、UNet 等）到您的仓库中的一个子文件夹。 - - -为了获得更好的结果，尝试使用更大的数据集进行更长时间的训练。我们只在较小规模的数据集上测试过此训练脚本。 - -
- -要使用 Weights and Biases 监控训练进度，请将 `--report_to=wandb` 参数添加到训练命令中，并使用 `--val_image_url` 指定验证图像，使用 `--validation_prompt` 指定验证提示。这对于调试模型非常有用。 - - +> [!TIP] +> 为了获得更好的结果，尝试使用更大的数据集进行更长时间的训练。我们只在较小规模的数据集上测试过此训练脚本。 +> +>
+> +> 要使用 Weights and Biases 监控训练进度，请将 `--report_to=wandb` 参数添加到训练命令中，并使用 `--val_image_url` 指定验证图像，使用 `--validation_prompt` 指定验证提示。这对于调试模型非常有用。 如果您在多个 GPU 上训练，请将 `--multi_gpu` 参数添加到 `accelerate launch` 命令中。 diff --git a/docs/source/zh/training/kandinsky.md b/docs/source/zh/training/kandinsky.md index 8da5c0c3a0..8ef3524ee7 100644 --- a/docs/source/zh/training/kandinsky.md +++ b/docs/source/zh/training/kandinsky.md @@ -9,11 +9,8 @@ http://www.apache.org/licenses/LICENSE-2.0 # Kandinsky 2.2 - - -此脚本是实验性的，容易过拟合并遇到灾难性遗忘等问题。尝试探索不同的超参数以在您的数据集上获得最佳结果。 - - +> [!WARNING] +> 此脚本是实验性的，容易过拟合并遇到灾难性遗忘等问题。尝试探索不同的超参数以在您的数据集上获得最佳结果。 Kandinsky 2.2 是一个多语言文本到图像模型，能够生成更逼真的图像。该模型包括一个图像先验模型，用于从文本提示创建图像嵌入，以及一个解码器模型，基于先验模型的嵌入生成图像。这就是为什么在 Diffusers 中您会找到两个独立的脚本用于 Kandinsky 2.2，一个用于训练先验模型，另一个用于训练解码器模型。您可以分别训练这两个模型，但为了获得最佳结果，您应该同时训练先验和解码器模型。 @@ -36,12 +33,9 @@ cd examples/kandinsky2_2/text_to_image pip install -r requirements.txt ``` - - -🤗 Accelerate 是一个帮助您在多个 GPU/TPU 上或使用混合精度进行训练的库。它会根据您的硬件和环境自动配置训练设置。查看 🤗 Accelerate 的 [快速入门](https://huggingface.co/docs/accelerate/quicktour -) 了解更多。 - - +> [!TIP] +> 🤗 Accelerate 是一个帮助您在多个 GPU/TPU 上或使用混合精度进行训练的库。它会根据您的硬件和环境自动配置训练设置。查看 🤗 Accelerate 的 [快速入门](https://huggingface.co/docs/accelerate/quicktour +> ) 了解更多。 初始化一个 🤗 Accelerate 环境： @@ -65,11 +59,8 @@ write_basic_config() 最后，如果您想在自己的数据集上训练模型，请查看 [创建用于训练的数据集](create_dataset) 指南，了解如何创建与训练脚本兼容的数据集。 - - -以下部分重点介绍了训练脚本中对于理解如何修改它很重要的部分，但并未详细涵盖脚本的每个方面。如果您有兴趣了解更多，请随时阅读脚本，并让我们知道您有任何疑问或顾虑。 - - +> [!TIP] +> 以下部分重点介绍了训练脚本中对于理解如何修改它很重要的部分，但并未详细涵盖脚本的每个方面。如果您有兴趣了解更多，请随时阅读脚本，并让我们知道您有任何疑问或顾虑。 ## 脚本参数 @@ -209,12 +200,9 @@ model_pred = unet(noisy_latents, timesteps, None, added_cond_kwargs=added_cond_k 如果您在多个GPU上训练，请在 `accelerate launch` 命令中添加 `--multi_gpu` 参数。 - - -要使用Weights & Biases监控训练进度，请在训练命令中添加 `--report_to=wandb` 参数。您还需要 -建议在训练命令中添加 `--validation_prompt` 以跟踪结果。这对于调试模型和查看中间结果非常有用。 - - +> [!TIP] +> 要使用Weights & Biases监控训练进度，请在训练命令中添加 `--report_to=wandb` 参数。您还需要 +> 建议在训练命令中添加 `--validation_prompt` 以跟踪结果。这对于调试模型和查看中间结果非常有用。 @@ -284,11 +272,8 @@ prompt="A robot naruto, 4k photo" image = pipeline(prompt=prompt, negative_prompt=negative_prompt).images[0] ``` - - -可以随意将 `kandinsky-community/kandinsky-2-2-decoder` 替换为您自己训练的 decoder 检查点！ - - +> [!TIP] +> 可以随意将 `kandinsky-community/kandinsky-2-2-decoder` 替换为您自己训练的 decoder 检查点！ diff --git a/docs/source/zh/training/lora.md b/docs/source/zh/training/lora.md index a7b7abb32d..ce29365450 100644 --- a/docs/source/zh/training/lora.md +++ b/docs/source/zh/training/lora.md @@ -12,19 +12,13 @@ specific language governing permissions and limitations under the License. # LoRA 低秩适配 - - -当前功能处于实验阶段，API可能在未来版本中变更。 - - +> [!WARNING] +> 当前功能处于实验阶段，API可能在未来版本中变更。 [LoRA（大语言模型的低秩适配）](https://hf.co/papers/2106.09685) 是一种轻量级训练技术，能显著减少可训练参数量。其原理是通过向模型注入少量新权重参数，仅训练这些新增参数。这使得LoRA训练速度更快、内存效率更高，并生成更小的模型权重文件（通常仅数百MB），便于存储和分享。LoRA还可与DreamBooth等其他训练技术结合以加速训练过程。 - - -LoRA具有高度通用性，目前已支持以下应用场景：[DreamBooth](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth_lora.py)、[Kandinsky 2.2](https://github.com/huggingface/diffusers/blob/main/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_decoder.py)、[Stable Diffusion XL](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_lora_sdxl.py)、[文生图](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_lora.py)以及[Wuerstchen](https://github.com/huggingface/diffusers/blob/main/examples/wuerstchen/text_to_image/train_text_to_image_lora_prior.py)。 - - +> [!TIP] +> LoRA具有高度通用性，目前已支持以下应用场景：[DreamBooth](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth_lora.py)、[Kandinsky 2.2](https://github.com/huggingface/diffusers/blob/main/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_decoder.py)、[Stable Diffusion XL](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_lora_sdxl.py)、[文生图](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_lora.py)以及[Wuerstchen](https://github.com/huggingface/diffusers/blob/main/examples/wuerstchen/text_to_image/train_text_to_image_lora_prior.py)。 本指南将通过解析[train_text_to_image_lora.py](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_lora.py)脚本，帮助您深入理解其工作原理，并掌握如何针对具体需求进行定制化修改。 @@ -57,11 +51,8 @@ pip install -r requirements_flax.txt - - -🤗 Accelerate是一个支持多GPU/TPU训练和混合精度计算的库，它能根据硬件环境自动配置训练方案。参阅🤗 Accelerate[快速入门](https://huggingface.co/docs/accelerate/quicktour)了解更多。 - - +> [!TIP] +> 🤗 Accelerate是一个支持多GPU/TPU训练和混合精度计算的库，它能根据硬件环境自动配置训练方案。参阅🤗 Accelerate[快速入门](https://huggingface.co/docs/accelerate/quicktour)了解更多。 初始化🤗 Accelerate环境： @@ -85,11 +76,8 @@ write_basic_config() 如需训练自定义数据集，请参考[创建训练数据集指南](create_dataset)了解数据准备流程。 - - -以下章节重点解析训练脚本中与LoRA相关的核心部分，但不会涵盖所有实现细节。如需完整理解，建议直接阅读[脚本源码](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_lora.py)，如有疑问欢迎反馈。 - - +> [!TIP] +> 以下章节重点解析训练脚本中与LoRA相关的核心部分，但不会涵盖所有实现细节。如需完整理解，建议直接阅读[脚本源码](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_lora.py)，如有疑问欢迎反馈。 ## 脚本参数 @@ -177,11 +165,8 @@ optimizer = optimizer_cls( 多GPU训练请添加`--multi_gpu`参数。 - - -在11GB显存的2080 Ti显卡上完整训练约需5小时。 - - +> [!WARNING] +> 在11GB显存的2080 Ti显卡上完整训练约需5小时。 ```bash export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5" diff --git a/docs/source/zh/training/text2image.md b/docs/source/zh/training/text2image.md index 193b839e9b..4465adbe2a 100644 --- a/docs/source/zh/training/text2image.md +++ b/docs/source/zh/training/text2image.md @@ -12,11 +12,8 @@ specific language governing permissions and limitations under the License. # 文生图 - - -文生图训练脚本目前处于实验阶段，容易出现过拟合和灾难性遗忘等问题。建议尝试不同超参数以获得最佳数据集适配效果。 - - +> [!WARNING] +> 文生图训练脚本目前处于实验阶段，容易出现过拟合和灾难性遗忘等问题。建议尝试不同超参数以获得最佳数据集适配效果。 Stable Diffusion 等文生图模型能够根据文本提示生成对应图像。 @@ -49,11 +46,8 @@ pip install -r requirements_flax.txt - - -🤗 Accelerate 是支持多GPU/TPU训练和混合精度的工具库，能根据硬件环境自动配置训练参数。参阅 🤗 Accelerate [快速入门](https://huggingface.co/docs/accelerate/quicktour) 了解更多。 - - +> [!TIP] +> 🤗 Accelerate 是支持多GPU/TPU训练和混合精度的工具库，能根据硬件环境自动配置训练参数。参阅 🤗 Accelerate [快速入门](https://huggingface.co/docs/accelerate/quicktour) 了解更多。 初始化 🤗 Accelerate 环境： @@ -79,11 +73,8 @@ write_basic_config() ## 脚本参数 - - -以下重点介绍脚本中影响训练效果的关键参数，如需完整参数说明可查阅 [脚本源码](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py)。如有疑问欢迎反馈。 - - +> [!TIP] +> 以下重点介绍脚本中影响训练效果的关键参数，如需完整参数说明可查阅 [脚本源码](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py)。如有疑问欢迎反馈。 训练脚本提供丰富参数供自定义训练流程，所有参数及说明详见 [`parse_args()`](https://github.com/huggingface/diffusers/blob/8959c5b9dec1c94d6ba482c94a58d2215c5fd026/examples/text_to_image/train_text_to_image.py#L193) 函数。该函数为每个参数提供默认值（如批次大小、学习率等），也可通过命令行参数覆盖。 @@ -160,11 +151,8 @@ def preprocess_train(examples): 以 [火影忍者BLIP标注数据集](https://huggingface.co/datasets/lambdalabs/naruto-blip-captions) 为例训练生成火影角色。设置环境变量 `MODEL_NAME` 和 `dataset_name` 指定模型和数据集（Hub或本地路径）。多GPU训练需在 `accelerate launch` 命令中添加 `--multi_gpu` 参数。 - - -使用本地数据集时，设置 `TRAIN_DIR` 和 `OUTPUT_DIR` 环境变量为数据集路径和模型保存路径。 - - +> [!TIP] +> 使用本地数据集时，设置 `TRAIN_DIR` 和 `OUTPUT_DIR` 环境变量为数据集路径和模型保存路径。 ```bash export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5" @@ -194,11 +182,8 @@ Flax训练方案在TPU/GPU上效率更高（由 [@duongna211](https://github.com 设置环境变量 `MODEL_NAME` 和 `dataset_name` 指定模型和数据集（Hub或本地路径）。 - - -使用本地数据集时，设置 `TRAIN_DIR` 和 `OUTPUT_DIR` 环境变量为数据集路径和模型保存路径。 - - +> [!TIP] +> 使用本地数据集时，设置 `TRAIN_DIR` 和 `OUTPUT_DIR` 环境变量为数据集路径和模型保存路径。 ```bash export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5" diff --git a/docs/source/zh/training/text_inversion.md b/docs/source/zh/training/text_inversion.md index 2945699c61..eda9f91144 100644 --- a/docs/source/zh/training/text_inversion.md +++ b/docs/source/zh/training/text_inversion.md @@ -45,11 +45,8 @@ pip install -r requirements_flax.txt - - -🤗 Accelerate 是一个帮助您在多GPU/TPU或混合精度环境下训练的工具库。它会根据硬件和环境自动配置训练设置。查看🤗 Accelerate [快速入门](https://huggingface.co/docs/accelerate/quicktour)了解更多。 - - +> [!TIP] +> 🤗 Accelerate 是一个帮助您在多GPU/TPU或混合精度环境下训练的工具库。它会根据硬件和环境自动配置训练设置。查看🤗 Accelerate [快速入门](https://huggingface.co/docs/accelerate/quicktour)了解更多。 初始化🤗 Accelerate环境： @@ -73,11 +70,8 @@ write_basic_config() 最后，如果想在自定义数据集上训练模型，请参阅[创建训练数据集](create_dataset)指南，了解如何创建适用于训练脚本的数据集。 - - -以下部分重点介绍训练脚本中需要理解的关键修改点，但未涵盖脚本所有细节。如需深入了解，可随时查阅[脚本源码](https://github.com/huggingface/diffusers/blob/main/examples/textual_inversion/textual_inversion.py)，如有疑问欢迎反馈。 - - +> [!TIP] +> 以下部分重点介绍训练脚本中需要理解的关键修改点，但未涵盖脚本所有细节。如需深入了解，可随时查阅[脚本源码](https://github.com/huggingface/diffusers/blob/main/examples/textual_inversion/textual_inversion.py)，如有疑问欢迎反馈。 ## 脚本参数 @@ -173,11 +167,8 @@ snapshot_download( - `token_identifier.txt`：特殊占位符词汇 - `type_of_concept.txt`：训练概念类型（"object"或"style"） - - -在单块V100 GPU上完整训练约需1小时。 - - +> [!WARNING] +> 在单块V100 GPU上完整训练约需1小时。 启动脚本前还有最后一步。如果想实时观察训练过程，可以定期保存生成图像。在训练命令中添加以下参数： diff --git a/docs/source/zh/training/wuerstchen.md b/docs/source/zh/training/wuerstchen.md index 8a6abe6624..c80cc944a3 100644 --- a/docs/source/zh/training/wuerstchen.md +++ b/docs/source/zh/training/wuerstchen.md @@ -33,11 +33,8 @@ cd examples/wuerstchen/text_to_image pip install -r requirements.txt ``` - - -🤗 Accelerate 是一个帮助您在多个 GPU/TPU 上或使用混合精度进行训练的库。它会根据您的硬件和环境自动配置训练设置。查看 🤗 Accelerate [快速入门](https://huggingface.co/docs/accelerate/quicktour) 以了解更多信息。 - - +> [!TIP] +> 🤗 Accelerate 是一个帮助您在多个 GPU/TPU 上或使用混合精度进行训练的库。它会根据您的硬件和环境自动配置训练设置。查看 🤗 Accelerate [快速入门](https://huggingface.co/docs/accelerate/quicktour) 以了解更多信息。 初始化一个 🤗 Accelerate 环境： @@ -61,11 +58,8 @@ write_basic_config() 最后，如果您想在自己的数据集上训练模型，请查看 [创建训练数据集](create_dataset) 指南，了解如何创建与训练脚本兼容的数据集。 - - -以下部分重点介绍了训练脚本中对于理解如何修改它很重要的部分，但并未涵盖 [脚本](https://github.com/huggingface/diffusers/blob/main/examples/wuerstchen/text_to_image/train_text_to_image_prior.py) 的详细信息。如果您有兴趣了解更多，请随时阅读脚本，并告诉我们您是否有任何问题或疑虑。 - - +> [!TIP] +> 以下部分重点介绍了训练脚本中对于理解如何修改它很重要的部分，但并未涵盖 [脚本](https://github.com/huggingface/diffusers/blob/main/examples/wuerstchen/text_to_image/train_text_to_image_prior.py) 的详细信息。如果您有兴趣了解更多，请随时阅读脚本，并告诉我们您是否有任何问题或疑虑。 ## 脚本参数 @@ -134,11 +128,8 @@ pred_noise = prior(noisy_latents, timesteps, prompt_embeds) 设置`DATASET_NAME`环境变量为Hub中的数据集名称。本指南使用[Naruto BLIP captions](https://huggingface.co/datasets/lambdalabs/naruto-blip-captions)数据集，但您也可以创建和训练自己的数据集（参见[创建用于训练的数据集](create_dataset)指南）。 - - -要使用Weights & Biases监控训练进度，请在训练命令中添加`--report_to=wandb`参数。您还需要在训练命令中添加`--validation_prompt`以跟踪结果。这对于调试模型和查看中间结果非常有用。 - - +> [!TIP] +> 要使用Weights & Biases监控训练进度，请在训练命令中添加`--report_to=wandb`参数。您还需要在训练命令中添加`--validation_prompt`以跟踪结果。这对于调试模型和查看中间结果非常有用。 ```bash export DATASET_NAME="lambdalabs/naruto-blip-captions" diff --git a/examples/community/matryoshka.py b/examples/community/matryoshka.py index 274851e2ac..3871552672 100644 --- a/examples/community/matryoshka.py +++ b/examples/community/matryoshka.py @@ -1475,11 +1475,8 @@ class MatryoshkaFusedAttnProcessor2_0: fused projection layers. For self-attention modules, all projection matrices (i.e., query, key, value) are fused. For cross-attention modules, key and value projection matrices are fused. - - - This API is currently 🧪 experimental in nature and can change in future. - - + > [!WARNING] + > This API is currently 🧪 experimental in nature and can change in future. """ def __init__(self): @@ -2696,11 +2693,8 @@ class MatryoshkaUNet2DConditionModel( Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) are fused. For cross-attention modules, key and value projection matrices are fused. - - - This API is 🧪 experimental. - - + > [!WARNING] + > This API is 🧪 experimental. """ self.original_attn_processors = None @@ -2719,11 +2713,8 @@ class MatryoshkaUNet2DConditionModel( def unfuse_qkv_projections(self): """Disables the fused QKV projection if enabled. - - - This API is 🧪 experimental. - - + > [!WARNING] + > This API is 🧪 experimental. """ if self.original_attn_processors is not None: diff --git a/examples/community/pipeline_stable_diffusion_boxdiff.py b/examples/community/pipeline_stable_diffusion_boxdiff.py index 1133321fcc..07e29b9c05 100644 --- a/examples/community/pipeline_stable_diffusion_boxdiff.py +++ b/examples/community/pipeline_stable_diffusion_boxdiff.py @@ -948,11 +948,8 @@ class StableDiffusionBoxDiffPipeline( Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) are fused. For cross-attention modules, key and value projection matrices are fused. - - - This API is 🧪 experimental. - - + > [!WARNING] + > This API is 🧪 experimental. Args: unet (`bool`, defaults to `True`): To apply fusion on the UNet. @@ -978,11 +975,8 @@ class StableDiffusionBoxDiffPipeline( def unfuse_qkv_projections(self, unet: bool = True, vae: bool = True): """Disable QKV projection fusion if enabled. - - - This API is 🧪 experimental. - - + > [!WARNING] + > This API is 🧪 experimental. Args: unet (`bool`, defaults to `True`): To apply fusion on the UNet. diff --git a/examples/community/pipeline_stable_diffusion_pag.py b/examples/community/pipeline_stable_diffusion_pag.py index 6728e2a60b..6b62b610af 100644 --- a/examples/community/pipeline_stable_diffusion_pag.py +++ b/examples/community/pipeline_stable_diffusion_pag.py @@ -940,9 +940,8 @@ class StableDiffusionPAGPipeline( """ Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) are fused. For cross-attention modules, key and value projection matrices are fused. - - This API is 🧪 experimental. - + > [!WARNING] + > This API is 🧪 experimental. Args: unet (`bool`, defaults to `True`): To apply fusion on the UNet. vae (`bool`, defaults to `True`): To apply fusion on the VAE. @@ -966,9 +965,8 @@ class StableDiffusionPAGPipeline( # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.unfuse_qkv_projections def unfuse_qkv_projections(self, unet: bool = True, vae: bool = True): """Disable QKV projection fusion if enabled. - - This API is 🧪 experimental. - + > [!WARNING] + > This API is 🧪 experimental. Args: unet (`bool`, defaults to `True`): To apply fusion on the UNet. vae (`bool`, defaults to `True`): To apply fusion on the VAE. diff --git a/examples/model_search/pipeline_easy.py b/examples/model_search/pipeline_easy.py index fcce297c37..ee5dced817 100644 --- a/examples/model_search/pipeline_easy.py +++ b/examples/model_search/pipeline_easy.py @@ -1246,12 +1246,9 @@ class EasyPipelineForText2Image(AutoPipelineForText2Image): Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when loading `from_flax`. - - - To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with - `hf auth login`. - - + > [!TIP] + > To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with + > `hf auth login`. Examples: @@ -1355,12 +1352,9 @@ class EasyPipelineForText2Image(AutoPipelineForText2Image): class). The overwritten components are passed directly to the pipelines `__init__` method. See example below for more information. - - - To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with - `hf auth login`. - - + > [!TIP] + > To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with + > `hf auth login`. Examples: @@ -1504,12 +1498,9 @@ class EasyPipelineForImage2Image(AutoPipelineForImage2Image): Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when loading `from_flax`. - - - To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with - `hf auth login`. - - + > [!TIP] + > To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with + > `hf auth login`. Examples: @@ -1614,12 +1605,9 @@ class EasyPipelineForImage2Image(AutoPipelineForImage2Image): class). The overwritten components are passed directly to the pipelines `__init__` method. See example below for more information. - - - To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with - `hf auth login`. - - + > [!TIP] + > To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with + > `hf auth login`. Examples: @@ -1763,12 +1751,9 @@ class EasyPipelineForInpainting(AutoPipelineForInpainting): Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when loading `from_flax`. - - - To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with - `hf auth login - - + > [!TIP] + > To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with + > `hf auth login Examples: @@ -1872,12 +1857,9 @@ class EasyPipelineForInpainting(AutoPipelineForInpainting): class). The overwritten components are passed directly to the pipelines `__init__` method. See example below for more information. - - - To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with - `hf auth login - - + > [!TIP] + > To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with + > `hf auth login Examples: diff --git a/src/diffusers/guiders/guider_utils.py b/src/diffusers/guiders/guider_utils.py index a6f2e76dc3..7524b5a3ea 100644 --- a/src/diffusers/guiders/guider_utils.py +++ b/src/diffusers/guiders/guider_utils.py @@ -247,15 +247,11 @@ class BaseGuidance(ConfigMixin, PushToHubMixin): The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier allowed by Git. - - - To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in with `hf - auth login`. You can also activate the special - ["offline-mode"](https://huggingface.co/diffusers/installation.html#offline-mode) to use this method in a + > [!TIP] > To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in + with `hf > auth login`. You can also activate the special > + ["offline-mode"](https://huggingface.co/diffusers/installation.html#offline-mode) to use this method in a > firewalled environment. - - """ config, kwargs, commit_hash = cls.load_config( pretrained_model_name_or_path=pretrained_model_name_or_path, diff --git a/src/diffusers/loaders/lora_base.py b/src/diffusers/loaders/lora_base.py index 0ee32f820b..3d75a7d875 100644 --- a/src/diffusers/loaders/lora_base.py +++ b/src/diffusers/loaders/lora_base.py @@ -544,11 +544,7 @@ class LoraBaseMixin: r""" Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - This is an experimental API. - - + > [!WARNING] > This is an experimental API. Args: components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into. @@ -628,11 +624,7 @@ class LoraBaseMixin: Reverses the effect of [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - This is an experimental API. - - + > [!WARNING] > This is an experimental API. Args: components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. diff --git a/src/diffusers/loaders/lora_pipeline.py b/src/diffusers/loaders/lora_pipeline.py index 65bdae6920..e25a29e1c0 100644 --- a/src/diffusers/loaders/lora_pipeline.py +++ b/src/diffusers/loaders/lora_pipeline.py @@ -246,13 +246,8 @@ class StableDiffusionLoraLoaderMixin(LoraBaseMixin): r""" Return state dict for lora weights and the network alphas. - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - + > [!WARNING] > We support loading A1111 formatted LoRA checkpoints in a limited capacity. > > This function is + experimental and might change in the future. Parameters: pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): @@ -545,11 +540,7 @@ class StableDiffusionLoraLoaderMixin(LoraBaseMixin): r""" Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - This is an experimental API. - - + > [!WARNING] > This is an experimental API. Args: components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into. @@ -586,11 +577,7 @@ class StableDiffusionLoraLoaderMixin(LoraBaseMixin): Reverses the effect of [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - This is an experimental API. - - + > [!WARNING] > This is an experimental API. Args: components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. @@ -698,13 +685,8 @@ class StableDiffusionXLLoraLoaderMixin(LoraBaseMixin): r""" Return state dict for lora weights and the network alphas. - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - + > [!WARNING] > We support loading A1111 formatted LoRA checkpoints in a limited capacity. > > This function is + experimental and might change in the future. Parameters: pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): @@ -2007,11 +1989,7 @@ class FluxLoraLoaderMixin(LoraBaseMixin): Reverses the effect of [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - This is an experimental API. - - + > [!WARNING] > This is an experimental API. Args: components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. diff --git a/src/diffusers/models/attention.py b/src/diffusers/models/attention.py index 3f59c8da8e..5164cf311d 100644 --- a/src/diffusers/models/attention.py +++ b/src/diffusers/models/attention.py @@ -111,11 +111,7 @@ class AttentionMixin: def unfuse_qkv_projections(self): """Disables the fused QKV projection if enabled. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ for module in self.modules(): if isinstance(module, AttentionModuleMixin): diff --git a/src/diffusers/models/attention_processor.py b/src/diffusers/models/attention_processor.py index 990245de17..66455d733a 100755 --- a/src/diffusers/models/attention_processor.py +++ b/src/diffusers/models/attention_processor.py @@ -3669,11 +3669,7 @@ class FusedAttnProcessor2_0: fused projection layers. For self-attention modules, all projection matrices (i.e., query, key, value) are fused. For cross-attention modules, key and value projection matrices are fused. - - - This API is currently 🧪 experimental in nature and can change in future. - - + > [!WARNING] > This API is currently 🧪 experimental in nature and can change in future. """ def __init__(self): diff --git a/src/diffusers/models/auto_model.py b/src/diffusers/models/auto_model.py index 47f3a992b3..a95b0ae64a 100644 --- a/src/diffusers/models/auto_model.py +++ b/src/diffusers/models/auto_model.py @@ -118,15 +118,11 @@ class AutoModel(ConfigMixin): trust_remote_cocde (`bool`, *optional*, defaults to `False`): Whether to trust remote code - - - To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in with `hf - auth login`. You can also activate the special - ["offline-mode"](https://huggingface.co/diffusers/installation.html#offline-mode) to use this method in a + > [!TIP] > To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in + with `hf > auth login`. You can also activate the special > + ["offline-mode"](https://huggingface.co/diffusers/installation.html#offline-mode) to use this method in a > firewalled environment. - - Example: ```py diff --git a/src/diffusers/models/autoencoders/autoencoder_kl.py b/src/diffusers/models/autoencoders/autoencoder_kl.py index 9a4375a36b..d823c2fb8b 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl.py @@ -532,11 +532,7 @@ class AutoencoderKL(ModelMixin, ConfigMixin, FromOriginalModelMixin, PeftAdapter Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) are fused. For cross-attention modules, key and value projection matrices are fused. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ self.original_attn_processors = None @@ -556,11 +552,7 @@ class AutoencoderKL(ModelMixin, ConfigMixin, FromOriginalModelMixin, PeftAdapter def unfuse_qkv_projections(self): """Disables the fused QKV projection if enabled. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ if self.original_attn_processors is not None: diff --git a/src/diffusers/models/controlnets/controlnet_sd3.py b/src/diffusers/models/controlnets/controlnet_sd3.py index 8d892cb3b6..0641c8bc01 100644 --- a/src/diffusers/models/controlnets/controlnet_sd3.py +++ b/src/diffusers/models/controlnets/controlnet_sd3.py @@ -270,11 +270,7 @@ class SD3ControlNetModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginal Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) are fused. For cross-attention modules, key and value projection matrices are fused. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ self.original_attn_processors = None @@ -294,11 +290,7 @@ class SD3ControlNetModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginal def unfuse_qkv_projections(self): """Disables the fused QKV projection if enabled. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ if self.original_attn_processors is not None: diff --git a/src/diffusers/models/controlnets/controlnet_xs.py b/src/diffusers/models/controlnets/controlnet_xs.py index aabae709e9..bcb4e25986 100644 --- a/src/diffusers/models/controlnets/controlnet_xs.py +++ b/src/diffusers/models/controlnets/controlnet_xs.py @@ -980,11 +980,7 @@ class UNetControlNetXSModel(ModelMixin, ConfigMixin): Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) are fused. For cross-attention modules, key and value projection matrices are fused. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ self.original_attn_processors = None @@ -1004,11 +1000,7 @@ class UNetControlNetXSModel(ModelMixin, ConfigMixin): def unfuse_qkv_projections(self): """Disables the fused QKV projection if enabled. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ if self.original_attn_processors is not None: diff --git a/src/diffusers/models/modeling_flax_utils.py b/src/diffusers/models/modeling_flax_utils.py index 8050afff27..fd19578321 100644 --- a/src/diffusers/models/modeling_flax_utils.py +++ b/src/diffusers/models/modeling_flax_utils.py @@ -227,15 +227,9 @@ class FlaxModelMixin(PushToHubMixin): This can be used to enable mixed-precision training or half-precision inference on GPUs or TPUs. If specified, all the computation will be performed with the given `dtype`. - - - This only specifies the dtype of the *computation* and does not influence the dtype of model - parameters. - - If you wish to change the dtype of the model parameters, see [`~FlaxModelMixin.to_fp16`] and - [`~FlaxModelMixin.to_bf16`]. - - + > [!TIP] > This only specifies the dtype of the *computation* and does not influence the dtype of model + > parameters. > > If you wish to change the dtype of the model parameters, see + [`~FlaxModelMixin.to_fp16`] and > [`~FlaxModelMixin.to_bf16`]. model_args (sequence of positional arguments, *optional*): All remaining positional arguments are passed to the underlying model's `__init__` method. diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index b3d74954bd..1af7ba9ac5 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -403,12 +403,8 @@ class ModelMixin(torch.nn.Module, PushToHubMixin): When this option is enabled, you should observe lower GPU memory usage and a potential speed up during inference. Speed up during training is not guaranteed. - - - ⚠️ When memory efficient attention and sliced attention are both enabled, memory efficient attention takes - precedent. - - + > [!WARNING] > ⚠️ When memory efficient attention and sliced attention are both enabled, memory efficient + attention takes > precedent. Parameters: attention_op (`Callable`, *optional*): @@ -917,15 +913,11 @@ class ModelMixin(torch.nn.Module, PushToHubMixin): Whether to disable mmap when loading a Safetensors model. This option can perform better when the model is on a network mount or hard drive, which may not handle the seeky-ness of mmap very well. - - - To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in with `hf - auth login`. You can also activate the special - ["offline-mode"](https://huggingface.co/diffusers/installation.html#offline-mode) to use this method in a + > [!TIP] > To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in + with `hf > auth login`. You can also activate the special > + ["offline-mode"](https://huggingface.co/diffusers/installation.html#offline-mode) to use this method in a > firewalled environment. - - Example: ```py diff --git a/src/diffusers/models/transformers/auraflow_transformer_2d.py b/src/diffusers/models/transformers/auraflow_transformer_2d.py index 4d7d1ba40e..bf6d9e1b38 100644 --- a/src/diffusers/models/transformers/auraflow_transformer_2d.py +++ b/src/diffusers/models/transformers/auraflow_transformer_2d.py @@ -431,11 +431,7 @@ class AuraFlowTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, From Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) are fused. For cross-attention modules, key and value projection matrices are fused. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ self.original_attn_processors = None @@ -455,11 +451,7 @@ class AuraFlowTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, From def unfuse_qkv_projections(self): """Disables the fused QKV projection if enabled. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ if self.original_attn_processors is not None: diff --git a/src/diffusers/models/transformers/cogvideox_transformer_3d.py b/src/diffusers/models/transformers/cogvideox_transformer_3d.py index 5038109690..9e0afdee66 100644 --- a/src/diffusers/models/transformers/cogvideox_transformer_3d.py +++ b/src/diffusers/models/transformers/cogvideox_transformer_3d.py @@ -397,11 +397,7 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, Cac Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) are fused. For cross-attention modules, key and value projection matrices are fused. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ self.original_attn_processors = None @@ -421,11 +417,7 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, Cac def unfuse_qkv_projections(self): """Disables the fused QKV projection if enabled. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ if self.original_attn_processors is not None: diff --git a/src/diffusers/models/transformers/hunyuan_transformer_2d.py b/src/diffusers/models/transformers/hunyuan_transformer_2d.py index f634718788..fbe9fe8df9 100644 --- a/src/diffusers/models/transformers/hunyuan_transformer_2d.py +++ b/src/diffusers/models/transformers/hunyuan_transformer_2d.py @@ -324,11 +324,7 @@ class HunyuanDiT2DModel(ModelMixin, ConfigMixin): Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) are fused. For cross-attention modules, key and value projection matrices are fused. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ self.original_attn_processors = None @@ -348,11 +344,7 @@ class HunyuanDiT2DModel(ModelMixin, ConfigMixin): def unfuse_qkv_projections(self): """Disables the fused QKV projection if enabled. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ if self.original_attn_processors is not None: diff --git a/src/diffusers/models/transformers/pixart_transformer_2d.py b/src/diffusers/models/transformers/pixart_transformer_2d.py index 40a14bfd9b..5a22144228 100644 --- a/src/diffusers/models/transformers/pixart_transformer_2d.py +++ b/src/diffusers/models/transformers/pixart_transformer_2d.py @@ -258,11 +258,7 @@ class PixArtTransformer2DModel(ModelMixin, ConfigMixin): Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) are fused. For cross-attention modules, key and value projection matrices are fused. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ self.original_attn_processors = None @@ -282,11 +278,7 @@ class PixArtTransformer2DModel(ModelMixin, ConfigMixin): def unfuse_qkv_projections(self): """Disables the fused QKV projection if enabled. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ if self.original_attn_processors is not None: diff --git a/src/diffusers/models/transformers/transformer_sd3.py b/src/diffusers/models/transformers/transformer_sd3.py index edf77a7df7..762d89c303 100644 --- a/src/diffusers/models/transformers/transformer_sd3.py +++ b/src/diffusers/models/transformers/transformer_sd3.py @@ -280,11 +280,7 @@ class SD3Transformer2DModel( Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) are fused. For cross-attention modules, key and value projection matrices are fused. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ self.original_attn_processors = None @@ -304,11 +300,7 @@ class SD3Transformer2DModel( def unfuse_qkv_projections(self): """Disables the fused QKV projection if enabled. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ if self.original_attn_processors is not None: diff --git a/src/diffusers/models/unets/unet_2d_condition.py b/src/diffusers/models/unets/unet_2d_condition.py index 736deb28c3..33bda8cb1e 100644 --- a/src/diffusers/models/unets/unet_2d_condition.py +++ b/src/diffusers/models/unets/unet_2d_condition.py @@ -872,11 +872,7 @@ class UNet2DConditionModel( Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) are fused. For cross-attention modules, key and value projection matrices are fused. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ self.original_attn_processors = None @@ -895,11 +891,7 @@ class UNet2DConditionModel( def unfuse_qkv_projections(self): """Disables the fused QKV projection if enabled. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ if self.original_attn_processors is not None: diff --git a/src/diffusers/models/unets/unet_3d_condition.py b/src/diffusers/models/unets/unet_3d_condition.py index bd67ea414a..b5151f3c9a 100644 --- a/src/diffusers/models/unets/unet_3d_condition.py +++ b/src/diffusers/models/unets/unet_3d_condition.py @@ -508,11 +508,7 @@ class UNet3DConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin) Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) are fused. For cross-attention modules, key and value projection matrices are fused. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ self.original_attn_processors = None @@ -532,11 +528,7 @@ class UNet3DConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin) def unfuse_qkv_projections(self): """Disables the fused QKV projection if enabled. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ if self.original_attn_processors is not None: diff --git a/src/diffusers/models/unets/unet_i2vgen_xl.py b/src/diffusers/models/unets/unet_i2vgen_xl.py index 8449bf894c..7148723a84 100644 --- a/src/diffusers/models/unets/unet_i2vgen_xl.py +++ b/src/diffusers/models/unets/unet_i2vgen_xl.py @@ -472,11 +472,7 @@ class I2VGenXLUNet(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin): Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) are fused. For cross-attention modules, key and value projection matrices are fused. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ self.original_attn_processors = None @@ -496,11 +492,7 @@ class I2VGenXLUNet(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin): def unfuse_qkv_projections(self): """Disables the fused QKV projection if enabled. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ if self.original_attn_processors is not None: diff --git a/src/diffusers/models/unets/unet_motion_model.py b/src/diffusers/models/unets/unet_motion_model.py index 0a112b5249..26616e53bd 100644 --- a/src/diffusers/models/unets/unet_motion_model.py +++ b/src/diffusers/models/unets/unet_motion_model.py @@ -1911,11 +1911,7 @@ class UNetMotionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin, Peft Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) are fused. For cross-attention modules, key and value projection matrices are fused. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ self.original_attn_processors = None @@ -1935,11 +1931,7 @@ class UNetMotionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin, Peft def unfuse_qkv_projections(self): """Disables the fused QKV projection if enabled. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ if self.original_attn_processors is not None: diff --git a/src/diffusers/modular_pipelines/components_manager.py b/src/diffusers/modular_pipelines/components_manager.py index ed847fa414..9dd8035c44 100644 --- a/src/diffusers/modular_pipelines/components_manager.py +++ b/src/diffusers/modular_pipelines/components_manager.py @@ -286,11 +286,7 @@ class ComponentsManager: encoders, etc.) across different modular pipelines. It includes features for duplicate detection, memory management, and component organization. - - - This is an experimental feature and is likely to change in the future. - - + > [!WARNING] > This is an experimental feature and is likely to change in the future. Example: ```python diff --git a/src/diffusers/modular_pipelines/flux/modular_pipeline.py b/src/diffusers/modular_pipelines/flux/modular_pipeline.py index 7d869041f2..563b033343 100644 --- a/src/diffusers/modular_pipelines/flux/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/flux/modular_pipeline.py @@ -25,11 +25,7 @@ class FluxModularPipeline(ModularPipeline, FluxLoraLoaderMixin, TextualInversion """ A ModularPipeline for Flux. - - - This is an experimental feature and is likely to change in the future. - - + > [!WARNING] > This is an experimental feature and is likely to change in the future. """ default_blocks_name = "FluxAutoBlocks" diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 206d19f173..037c9e323c 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -226,11 +226,7 @@ class ModularPipelineBlocks(ConfigMixin, PushToHubMixin): [`ModularPipelineBlocks`] provides method to load and save the definition of pipeline blocks. - - - This is an experimental feature and is likely to change in the future. - - + > [!WARNING] > This is an experimental feature and is likely to change in the future. """ config_name = "modular_config.json" @@ -525,11 +521,7 @@ class AutoPipelineBlocks(ModularPipelineBlocks): This class inherits from [`ModularPipelineBlocks`]. Check the superclass documentation for the generic methods the library implements for all the pipeline blocks (such as loading or saving etc.) - - - This is an experimental feature and is likely to change in the future. - - + > [!WARNING] > This is an experimental feature and is likely to change in the future. Attributes: block_classes: List of block classes to be used @@ -787,11 +779,7 @@ class SequentialPipelineBlocks(ModularPipelineBlocks): This class inherits from [`ModularPipelineBlocks`]. Check the superclass documentation for the generic methods the library implements for all the pipeline blocks (such as loading or saving etc.) - - - This is an experimental feature and is likely to change in the future. - - + > [!WARNING] > This is an experimental feature and is likely to change in the future. Attributes: block_classes: List of block classes to be used @@ -1146,11 +1134,7 @@ class LoopSequentialPipelineBlocks(ModularPipelineBlocks): This class inherits from [`ModularPipelineBlocks`]. Check the superclass documentation for the generic methods the library implements for all the pipeline blocks (such as loading or saving etc.) - - - This is an experimental feature and is likely to change in the future. - - + > [!WARNING] > This is an experimental feature and is likely to change in the future. Attributes: block_classes: List of block classes to be used @@ -1433,11 +1417,7 @@ class ModularPipeline(ConfigMixin, PushToHubMixin): """ Base class for all Modular pipelines. - - - This is an experimental feature and is likely to change in the future. - - + > [!WARNING] > This is an experimental feature and is likely to change in the future. Args: blocks: ModularPipelineBlocks, the blocks to be used in the pipeline @@ -2173,12 +2153,8 @@ class ModularPipeline(ConfigMixin, PushToHubMixin): Performs Pipeline dtype and/or device conversion. A torch.dtype and torch.device are inferred from the arguments of `self.to(*args, **kwargs).` - - - If the pipeline already has the correct torch.dtype and torch.device, then it is returned as is. Otherwise, - the returned pipeline is a copy of self with the desired torch.dtype and torch.device. - - + > [!TIP] > If the pipeline already has the correct torch.dtype and torch.device, then it is returned as is. + Otherwise, > the returned pipeline is a copy of self with the desired torch.dtype and torch.device. Here are the ways to call `to`: diff --git a/src/diffusers/modular_pipelines/node_utils.py b/src/diffusers/modular_pipelines/node_utils.py new file mode 100644 index 0000000000..f7ee1dd309 --- /dev/null +++ b/src/diffusers/modular_pipelines/node_utils.py @@ -0,0 +1,661 @@ +import json +import logging +import os +from pathlib import Path +from typing import List, Optional, Tuple, Union + +import numpy as np +import PIL +import torch + +from ..configuration_utils import ConfigMixin +from ..image_processor import PipelineImageInput +from .modular_pipeline import ModularPipelineBlocks, SequentialPipelineBlocks +from .modular_pipeline_utils import InputParam + + +logger = logging.getLogger(__name__) + +# YiYi Notes: this is actually for SDXL, put it here for now +SDXL_INPUTS_SCHEMA = { + "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", + ), + "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" + ), + "eta": InputParam("eta", type_hint=float, default=0.0, description="Parameter η in the DDIM paper"), + "output_type": InputParam( + "output_type", type_hint=str, default="pil", description="Output format (pil/tensor/np.array)" + ), + "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_SCHEMA = { + "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_PARAM_SCHEMA = {**SDXL_INPUTS_SCHEMA, **SDXL_INTERMEDIATE_INPUTS_SCHEMA} + + +DEFAULT_PARAM_MAPS = { + "prompt": { + "label": "Prompt", + "type": "string", + "default": "a bear sitting in a chair drinking a milkshake", + "display": "textarea", + }, + "negative_prompt": { + "label": "Negative Prompt", + "type": "string", + "default": "deformed, ugly, wrong proportion, low res, bad anatomy, worst quality, low quality", + "display": "textarea", + }, + "num_inference_steps": { + "label": "Steps", + "type": "int", + "default": 25, + "min": 1, + "max": 1000, + }, + "seed": { + "label": "Seed", + "type": "int", + "default": 0, + "min": 0, + "display": "random", + }, + "width": { + "label": "Width", + "type": "int", + "display": "text", + "default": 1024, + "min": 8, + "max": 8192, + "step": 8, + "group": "dimensions", + }, + "height": { + "label": "Height", + "type": "int", + "display": "text", + "default": 1024, + "min": 8, + "max": 8192, + "step": 8, + "group": "dimensions", + }, + "images": { + "label": "Images", + "type": "image", + "display": "output", + }, + "image": { + "label": "Image", + "type": "image", + "display": "input", + }, +} + +DEFAULT_TYPE_MAPS = { + "int": { + "type": "int", + "default": 0, + "min": 0, + }, + "float": { + "type": "float", + "default": 0.0, + "min": 0.0, + }, + "str": { + "type": "string", + "default": "", + }, + "bool": { + "type": "boolean", + "default": False, + }, + "image": { + "type": "image", + }, +} + +DEFAULT_MODEL_KEYS = ["unet", "vae", "text_encoder", "tokenizer", "controlnet", "transformer", "image_encoder"] +DEFAULT_CATEGORY = "Modular Diffusers" +DEFAULT_EXCLUDE_MODEL_KEYS = ["processor", "feature_extractor", "safety_checker"] +DEFAULT_PARAMS_GROUPS_KEYS = { + "text_encoders": ["text_encoder", "tokenizer"], + "ip_adapter_embeds": ["ip_adapter_embeds"], + "prompt_embeddings": ["prompt_embeds"], +} + + +def get_group_name(name, group_params_keys=DEFAULT_PARAMS_GROUPS_KEYS): + """ + Get the group name for a given parameter name, if not part of a group, return None e.g. "prompt_embeds" -> + "text_embeds", "text_encoder" -> "text_encoders", "prompt" -> None + """ + if name is None: + return None + for group_name, group_keys in group_params_keys.items(): + for group_key in group_keys: + if group_key in name: + return group_name + return None + + +class ModularNode(ConfigMixin): + """ + A ModularNode is a base class to build UI nodes using diffusers. Currently only supports Mellon. It is a wrapper + around a ModularPipelineBlocks object. + + > [!WARNING] > This is an experimental feature and is likely to change in the future. + """ + + config_name = "node_config.json" + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: str, + trust_remote_code: Optional[bool] = None, + **kwargs, + ): + blocks = ModularPipelineBlocks.from_pretrained( + pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **kwargs + ) + return cls(blocks, **kwargs) + + def __init__(self, blocks, category=DEFAULT_CATEGORY, label=None, **kwargs): + self.blocks = blocks + + if label is None: + label = self.blocks.__class__.__name__ + # blocks param name -> mellon param name + self.name_mapping = {} + + input_params = {} + # pass or create a default param dict for each input + # e.g. for prompt, + # prompt = { + # "name": "text_input", # the name of the input in node definition, could be different from the input name in diffusers + # "label": "Prompt", + # "type": "string", + # "default": "a bear sitting in a chair drinking a milkshake", + # "display": "textarea"} + # if type is not specified, it'll be a "custom" param of its own type + # e.g. you can pass ModularNode(scheduler = {name :"scheduler"}) + # it will get this spec in node definition {"scheduler": {"label": "Scheduler", "type": "scheduler", "display": "input"}} + # name can be a dict, in that case, it is part of a "dict" input in mellon nodes, e.g. text_encoder= {name: {"text_encoders": "text_encoder"}} + inputs = self.blocks.inputs + self.blocks.intermediate_inputs + for inp in inputs: + param = kwargs.pop(inp.name, None) + if param: + # user can pass a param dict for all inputs, e.g. ModularNode(prompt = {...}) + input_params[inp.name] = param + mellon_name = param.pop("name", inp.name) + if mellon_name != inp.name: + self.name_mapping[inp.name] = mellon_name + continue + + if inp.name not in DEFAULT_PARAM_MAPS and not inp.required and not get_group_name(inp.name): + continue + + if inp.name in DEFAULT_PARAM_MAPS: + # first check if it's in the default param map, if so, directly use that + param = DEFAULT_PARAM_MAPS[inp.name].copy() + elif get_group_name(inp.name): + param = get_group_name(inp.name) + if inp.name not in self.name_mapping: + self.name_mapping[inp.name] = param + else: + # if not, check if it's in the SDXL input schema, if so, + # 1. use the type hint to determine the type + # 2. use the default param dict for the type e.g. if "steps" is a "int" type, {"steps": {"type": "int", "default": 0, "min": 0}} + if inp.type_hint is not None: + type_str = str(inp.type_hint).lower() + else: + inp_spec = SDXL_PARAM_SCHEMA.get(inp.name, None) + type_str = str(inp_spec.type_hint).lower() if inp_spec else "" + for type_key, type_param in DEFAULT_TYPE_MAPS.items(): + if type_key in type_str: + param = type_param.copy() + param["label"] = inp.name + param["display"] = "input" + break + else: + param = inp.name + # add the param dict to the inp_params dict + input_params[inp.name] = param + + component_params = {} + for comp in self.blocks.expected_components: + param = kwargs.pop(comp.name, None) + if param: + component_params[comp.name] = param + mellon_name = param.pop("name", comp.name) + if mellon_name != comp.name: + self.name_mapping[comp.name] = mellon_name + continue + + to_exclude = False + for exclude_key in DEFAULT_EXCLUDE_MODEL_KEYS: + if exclude_key in comp.name: + to_exclude = True + break + if to_exclude: + continue + + if get_group_name(comp.name): + param = get_group_name(comp.name) + if comp.name not in self.name_mapping: + self.name_mapping[comp.name] = param + elif comp.name in DEFAULT_MODEL_KEYS: + param = {"label": comp.name, "type": "diffusers_auto_model", "display": "input"} + else: + param = comp.name + # add the param dict to the model_params dict + component_params[comp.name] = param + + output_params = {} + if isinstance(self.blocks, SequentialPipelineBlocks): + last_block_name = list(self.blocks.sub_blocks.keys())[-1] + outputs = self.blocks.sub_blocks[last_block_name].intermediate_outputs + else: + outputs = self.blocks.intermediate_outputs + + for out in outputs: + param = kwargs.pop(out.name, None) + if param: + output_params[out.name] = param + mellon_name = param.pop("name", out.name) + if mellon_name != out.name: + self.name_mapping[out.name] = mellon_name + continue + + if out.name in DEFAULT_PARAM_MAPS: + param = DEFAULT_PARAM_MAPS[out.name].copy() + param["display"] = "output" + else: + group_name = get_group_name(out.name) + if group_name: + param = group_name + if out.name not in self.name_mapping: + self.name_mapping[out.name] = param + else: + param = out.name + # add the param dict to the outputs dict + output_params[out.name] = param + + if len(kwargs) > 0: + logger.warning(f"Unused kwargs: {kwargs}") + + register_dict = { + "category": category, + "label": label, + "input_params": input_params, + "component_params": component_params, + "output_params": output_params, + "name_mapping": self.name_mapping, + } + self.register_to_config(**register_dict) + + def setup(self, components_manager, collection=None): + self.pipeline = self.blocks.init_pipeline(components_manager=components_manager, collection=collection) + self._components_manager = components_manager + + @property + def mellon_config(self): + return self._convert_to_mellon_config() + + def _convert_to_mellon_config(self): + node = {} + node["label"] = self.config.label + node["category"] = self.config.category + + node_param = {} + for inp_name, inp_param in self.config.input_params.items(): + if inp_name in self.name_mapping: + mellon_name = self.name_mapping[inp_name] + else: + mellon_name = inp_name + if isinstance(inp_param, str): + param = { + "label": inp_param, + "type": inp_param, + "display": "input", + } + else: + param = inp_param + + if mellon_name not in node_param: + node_param[mellon_name] = param + else: + logger.debug(f"Input param {mellon_name} already exists in node_param, skipping {inp_name}") + + for comp_name, comp_param in self.config.component_params.items(): + if comp_name in self.name_mapping: + mellon_name = self.name_mapping[comp_name] + else: + mellon_name = comp_name + if isinstance(comp_param, str): + param = { + "label": comp_param, + "type": comp_param, + "display": "input", + } + else: + param = comp_param + + if mellon_name not in node_param: + node_param[mellon_name] = param + else: + logger.debug(f"Component param {comp_param} already exists in node_param, skipping {comp_name}") + + for out_name, out_param in self.config.output_params.items(): + if out_name in self.name_mapping: + mellon_name = self.name_mapping[out_name] + else: + mellon_name = out_name + if isinstance(out_param, str): + param = { + "label": out_param, + "type": out_param, + "display": "output", + } + else: + param = out_param + + if mellon_name not in node_param: + node_param[mellon_name] = param + else: + logger.debug(f"Output param {out_param} already exists in node_param, skipping {out_name}") + node["params"] = node_param + return node + + def save_mellon_config(self, file_path): + """ + Save the Mellon configuration to a JSON file. + + Args: + file_path (str or Path): Path where the JSON file will be saved + + Returns: + Path: Path to the saved config file + """ + file_path = Path(file_path) + + # Create directory if it doesn't exist + os.makedirs(file_path.parent, exist_ok=True) + + # Create a combined dictionary with module definition and name mapping + config = {"module": self.mellon_config, "name_mapping": self.name_mapping} + + # Save the config to file + with open(file_path, "w", encoding="utf-8") as f: + json.dump(config, f, indent=2) + + logger.info(f"Mellon config and name mapping saved to {file_path}") + + return file_path + + @classmethod + def load_mellon_config(cls, file_path): + """ + Load a Mellon configuration from a JSON file. + + Args: + file_path (str or Path): Path to the JSON file containing Mellon config + + Returns: + dict: The loaded combined configuration containing 'module' and 'name_mapping' + """ + file_path = Path(file_path) + + if not file_path.exists(): + raise FileNotFoundError(f"Config file not found: {file_path}") + + with open(file_path, "r", encoding="utf-8") as f: + config = json.load(f) + + logger.info(f"Mellon config loaded from {file_path}") + + return config + + def process_inputs(self, **kwargs): + params_components = {} + for comp_name, comp_param in self.config.component_params.items(): + logger.debug(f"component: {comp_name}") + mellon_comp_name = self.name_mapping.get(comp_name, comp_name) + if mellon_comp_name in kwargs: + if isinstance(kwargs[mellon_comp_name], dict) and comp_name in kwargs[mellon_comp_name]: + comp = kwargs[mellon_comp_name].pop(comp_name) + else: + comp = kwargs.pop(mellon_comp_name) + if comp: + params_components[comp_name] = self._components_manager.get_one(comp["model_id"]) + + params_run = {} + for inp_name, inp_param in self.config.input_params.items(): + logger.debug(f"input: {inp_name}") + mellon_inp_name = self.name_mapping.get(inp_name, inp_name) + if mellon_inp_name in kwargs: + if isinstance(kwargs[mellon_inp_name], dict) and inp_name in kwargs[mellon_inp_name]: + inp = kwargs[mellon_inp_name].pop(inp_name) + else: + inp = kwargs.pop(mellon_inp_name) + if inp is not None: + params_run[inp_name] = inp + + return_output_names = list(self.config.output_params.keys()) + + return params_components, params_run, return_output_names + + def execute(self, **kwargs): + params_components, params_run, return_output_names = self.process_inputs(**kwargs) + + self.pipeline.update_components(**params_components) + output = self.pipeline(**params_run, output=return_output_names) + return output diff --git a/src/diffusers/modular_pipelines/qwenimage/modular_pipeline.py b/src/diffusers/modular_pipelines/qwenimage/modular_pipeline.py index 3248d13159..7200169923 100644 --- a/src/diffusers/modular_pipelines/qwenimage/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/qwenimage/modular_pipeline.py @@ -97,11 +97,7 @@ class QwenImageModularPipeline(ModularPipeline, QwenImageLoraLoaderMixin): """ A ModularPipeline for QwenImage. - - - This is an experimental feature and is likely to change in the future. - - + > [!WARNING] > This is an experimental feature and is likely to change in the future. """ default_blocks_name = "QwenImageAutoBlocks" @@ -153,11 +149,7 @@ class QwenImageEditModularPipeline(ModularPipeline, QwenImageLoraLoaderMixin): """ A ModularPipeline for QwenImage-Edit. - - - This is an experimental feature and is likely to change in the future. - - + > [!WARNING] > This is an experimental feature and is likely to change in the future. """ default_blocks_name = "QwenImageEditAutoBlocks" diff --git a/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_pipeline.py b/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_pipeline.py index 29a717f72e..f2a4c96073 100644 --- a/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_pipeline.py @@ -47,11 +47,7 @@ class StableDiffusionXLModularPipeline( """ A ModularPipeline for Stable Diffusion XL. - - - This is an experimental feature and is likely to change in the future. - - + > [!WARNING] > This is an experimental feature and is likely to change in the future. """ default_blocks_name = "StableDiffusionXLAutoBlocks" diff --git a/src/diffusers/modular_pipelines/wan/modular_pipeline.py b/src/diffusers/modular_pipelines/wan/modular_pipeline.py index da4aada438..e4adf3d151 100644 --- a/src/diffusers/modular_pipelines/wan/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/wan/modular_pipeline.py @@ -30,11 +30,7 @@ class WanModularPipeline( """ A ModularPipeline for Wan. - - - This is an experimental feature and is likely to change in the future. - - + > [!WARNING] > This is an experimental feature and is likely to change in the future. """ default_blocks_name = "WanAutoBlocks" diff --git a/src/diffusers/pipelines/auto_pipeline.py b/src/diffusers/pipelines/auto_pipeline.py index 880984eeb8..d265bfdcaf 100644 --- a/src/diffusers/pipelines/auto_pipeline.py +++ b/src/diffusers/pipelines/auto_pipeline.py @@ -407,12 +407,8 @@ class AutoPipelineForText2Image(ConfigMixin): Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when loading `from_flax`. - - - To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with `hf - auth login`. - - + > [!TIP] > To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in + with `hf > auth login`. Examples: @@ -702,12 +698,8 @@ class AutoPipelineForImage2Image(ConfigMixin): Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when loading `from_flax`. - - - To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with `hf - auth login`. - - + > [!TIP] > To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in + with `hf > auth login`. Examples: @@ -1012,12 +1004,8 @@ class AutoPipelineForInpainting(ConfigMixin): Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when loading `from_flax`. - - - To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with `hf - auth login`. - - + > [!TIP] > To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in + with `hf > auth login`. Examples: diff --git a/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py b/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py index 41303d9c5c..6de8e5747b 100644 --- a/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py +++ b/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py @@ -146,16 +146,13 @@ class StableDiffusionControlNetInpaintPipeline( - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters - - - This pipeline can be used with checkpoints that have been specifically fine-tuned for inpainting + > [!TIP] > This pipeline can be used with checkpoints that have been specifically fine-tuned for inpainting > ([stable-diffusion-v1-5/stable-diffusion-inpainting](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-inpainting)) - as well as default text-to-image Stable Diffusion checkpoints + > as well as default text-to-image Stable Diffusion checkpoints > ([stable-diffusion-v1-5/stable-diffusion-v1-5](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5)). - Default text-to-image Stable Diffusion checkpoints might be preferable for ControlNets that have been fine-tuned on - those, such as [lllyasviel/control_v11p_sd15_inpaint](https://huggingface.co/lllyasviel/control_v11p_sd15_inpaint). - - + > Default text-to-image Stable Diffusion checkpoints might be preferable for ControlNets that have been fine-tuned + on > those, such as + [lllyasviel/control_v11p_sd15_inpaint](https://huggingface.co/lllyasviel/control_v11p_sd15_inpaint). Args: vae ([`AutoencoderKL`]): diff --git a/src/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py b/src/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py index 1de1d4bde7..d4c6f336df 100644 --- a/src/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py +++ b/src/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py @@ -394,12 +394,8 @@ class FlaxStableDiffusionControlNetPipeline(FlaxDiffusionPipeline): jit (`bool`, defaults to `False`): Whether to run `pmap` versions of the generation and safety scoring functions. - - - This argument exists because `__call__` is not yet end-to-end pmap-able. It will be removed in a - future release. - - + > [!WARNING] > This argument exists because `__call__` is not yet end-to-end pmap-able. It will be + removed in a > future release. Examples: diff --git a/src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py b/src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py index eda950998d..397fbc0d85 100644 --- a/src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py +++ b/src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py @@ -1000,11 +1000,7 @@ class UNetFlatConditionModel(ModelMixin, ConfigMixin): Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) are fused. For cross-attention modules, key and value projection matrices are fused. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ self.original_attn_processors = None @@ -1021,11 +1017,7 @@ class UNetFlatConditionModel(ModelMixin, ConfigMixin): def unfuse_qkv_projections(self): """Disables the fused QKV projection if enabled. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. """ if self.original_attn_processors is not None: diff --git a/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_inpaint.py b/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_inpaint.py index 4c02b3dd6d..3daaac328c 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_inpaint.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_inpaint.py @@ -150,17 +150,13 @@ class StableDiffusionControlNetPAGInpaintPipeline( - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters - - - This pipeline can be used with checkpoints that have been specifically fine-tuned for inpainting - ([runwayml/stable-diffusion-inpainting](https://huggingface.co/runwayml/stable-diffusion-inpainting)) as well as - default text-to-image Stable Diffusion checkpoints - ([runwayml/stable-diffusion-v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5)). Default text-to-image - Stable Diffusion checkpoints might be preferable for ControlNets that have been fine-tuned on those, such as + > [!TIP] > This pipeline can be used with checkpoints that have been specifically fine-tuned for inpainting > + ([runwayml/stable-diffusion-inpainting](https://huggingface.co/runwayml/stable-diffusion-inpainting)) as well as > + default text-to-image Stable Diffusion checkpoints > + ([runwayml/stable-diffusion-v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5)). Default text-to-image > + Stable Diffusion checkpoints might be preferable for ControlNets that have been fine-tuned on those, such as > [lllyasviel/control_v11p_sd15_inpaint](https://huggingface.co/lllyasviel/control_v11p_sd15_inpaint). - - Args: vae ([`AutoencoderKL`]): Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. diff --git a/src/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py b/src/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py index 3e22c9a845..61435b80ca 100644 --- a/src/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py +++ b/src/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py @@ -158,11 +158,7 @@ def prepare_mask_and_masked_image(image, mask): class PaintByExamplePipeline(DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin): _last_supported_version = "0.33.1" r""" - - - 🧪 This is an experimental feature! - - + > [!WARNING] > 🧪 This is an experimental feature! Pipeline for image-guided image inpainting using Stable Diffusion. diff --git a/src/diffusers/pipelines/pipeline_flax_utils.py b/src/diffusers/pipelines/pipeline_flax_utils.py index f69968022e..2724c764c7 100644 --- a/src/diffusers/pipelines/pipeline_flax_utils.py +++ b/src/diffusers/pipelines/pipeline_flax_utils.py @@ -276,12 +276,8 @@ class FlaxDiffusionPipeline(ConfigMixin, PushToHubMixin): Can be used to overwrite load and saveable variables (the pipeline components) of the specific pipeline class. The overwritten components are passed directly to the pipelines `__init__` method. - - - To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in with `hf - auth login`. - - + > [!TIP] > To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in + with `hf > auth login`. Examples: diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index 3f6e53099b..392d5fb3fe 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -372,12 +372,8 @@ class DiffusionPipeline(ConfigMixin, PushToHubMixin): Performs Pipeline dtype and/or device conversion. A torch.dtype and torch.device are inferred from the arguments of `self.to(*args, **kwargs).` - - - If the pipeline already has the correct torch.dtype and torch.device, then it is returned as is. Otherwise, - the returned pipeline is a copy of self with the desired torch.dtype and torch.device. - - + > [!TIP] > If the pipeline already has the correct torch.dtype and torch.device, then it is returned as is. + Otherwise, > the returned pipeline is a copy of self with the desired torch.dtype and torch.device. Here are the ways to call `to`: @@ -627,11 +623,7 @@ class DiffusionPipeline(ConfigMixin, PushToHubMixin): `torch.float32` is used. custom_pipeline (`str`, *optional*): - - - 🧪 This is an experimental feature and may change in the future. - - + > [!WARNING] > 🧪 This is an experimental feature and may change in the future. Can be either: @@ -716,12 +708,8 @@ class DiffusionPipeline(ConfigMixin, PushToHubMixin): dduf_file(`str`, *optional*): Load weights from the specified dduf file. - - - To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with `hf - auth login`. - - + > [!TIP] > To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in + with `hf > auth login`. Examples: @@ -1508,11 +1496,7 @@ class DiffusionPipeline(ConfigMixin, PushToHubMixin): - A path to a *directory* (`./my_pipeline_directory/`) containing a custom pipeline. The directory must contain a file called `pipeline.py` that defines the custom pipeline. - - - 🧪 This is an experimental feature and may change in the future. - - + > [!WARNING] > 🧪 This is an experimental feature and may change in the future. For more information on how to load and create custom pipelines, take a look at [How to contribute a community pipeline](https://huggingface.co/docs/diffusers/main/en/using-diffusers/contribute_pipeline). @@ -1566,12 +1550,8 @@ class DiffusionPipeline(ConfigMixin, PushToHubMixin): `os.PathLike`: A path to the downloaded pipeline. - - - To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in with `hf - auth login - - + > [!TIP] > To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in + with `hf > auth login """ cache_dir = kwargs.pop("cache_dir", None) @@ -1944,12 +1924,8 @@ class DiffusionPipeline(ConfigMixin, PushToHubMixin): option is enabled, you should observe lower GPU memory usage and a potential speed up during inference. Speed up during training is not guaranteed. - - - ⚠️ When memory efficient attention and sliced attention are both enabled, memory efficient attention takes - precedent. - - + > [!WARNING] > ⚠️ When memory efficient attention and sliced attention are both enabled, memory efficient + attention takes > precedent. Parameters: attention_op (`Callable`, *optional*): @@ -2005,13 +1981,10 @@ class DiffusionPipeline(ConfigMixin, PushToHubMixin): in slices to compute attention in several steps. For more than one attention head, the computation is performed sequentially over each head. This is useful to save some memory in exchange for a small speed decrease. - - - ⚠️ Don't enable attention slicing if you're already using `scaled_dot_product_attention` (SDPA) from PyTorch - 2.0 or xFormers. These attention computations are already very memory efficient so you won't need to enable - this function. If you enable attention slicing with SDPA or xFormers, it can lead to serious slow downs! - - + > [!WARNING] > ⚠️ Don't enable attention slicing if you're already using `scaled_dot_product_attention` (SDPA) + from PyTorch > 2.0 or xFormers. These attention computations are already very memory efficient so you won't + need to enable > this function. If you enable attention slicing with SDPA or xFormers, it can lead to serious + slow downs! Args: slice_size (`str` or `int`, *optional*, defaults to `"auto"`): @@ -2288,11 +2261,7 @@ class StableDiffusionMixin: Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) are fused. For cross-attention modules, key and value projection matrices are fused. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. Args: unet (`bool`, defaults to `True`): To apply fusion on the UNet. @@ -2317,11 +2286,7 @@ class StableDiffusionMixin: def unfuse_qkv_projections(self, unet: bool = True, vae: bool = True): """Disable QKV projection fusion if enabled. - - - This API is 🧪 experimental. - - + > [!WARNING] > This API is 🧪 experimental. Args: unet (`bool`, defaults to `True`): To apply fusion on the UNet. diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py index 1afa7698da..6befe77aa4 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py @@ -349,12 +349,8 @@ class FlaxStableDiffusionPipeline(FlaxDiffusionPipeline): jit (`bool`, defaults to `False`): Whether to run `pmap` versions of the generation and safety scoring functions. - - - This argument exists because `__call__` is not yet end-to-end pmap-able. It will be removed in a - future release. - - + > [!WARNING] > This argument exists because `__call__` is not yet end-to-end pmap-able. It will be + removed in a > future release. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] instead of diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py index 78e3ba239c..81656beba7 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py @@ -389,12 +389,8 @@ class FlaxStableDiffusionImg2ImgPipeline(FlaxDiffusionPipeline): jit (`bool`, defaults to `False`): Whether to run `pmap` versions of the generation and safety scoring functions. - - - This argument exists because `__call__` is not yet end-to-end pmap-able. It will be removed in a - future release. - - + > [!WARNING] > This argument exists because `__call__` is not yet end-to-end pmap-able. It will be + removed in a > future release. Examples: diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py index b7e17ba681..5938fe232a 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py @@ -103,11 +103,7 @@ class FlaxStableDiffusionInpaintPipeline(FlaxDiffusionPipeline): r""" Flax-based pipeline for text-guided image inpainting using Stable Diffusion. - - - 🧪 This is an experimental feature! - - + > [!WARNING] > 🧪 This is an experimental feature! This model inherits from [`FlaxDiffusionPipeline`]. Check the superclass documentation for the generic methods implemented for all pipelines (downloading, saving, running on a particular device, etc.). @@ -435,12 +431,8 @@ class FlaxStableDiffusionInpaintPipeline(FlaxDiffusionPipeline): jit (`bool`, defaults to `False`): Whether to run `pmap` versions of the generation and safety scoring functions. - - - This argument exists because `__call__` is not yet end-to-end pmap-able. It will be removed in a - future release. - - + > [!WARNING] > This argument exists because `__call__` is not yet end-to-end pmap-able. It will be + removed in a > future release. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] instead of diff --git a/src/diffusers/pipelines/stable_diffusion_diffedit/pipeline_stable_diffusion_diffedit.py b/src/diffusers/pipelines/stable_diffusion_diffedit/pipeline_stable_diffusion_diffedit.py index 87bd9f4444..65c25ffbe4 100644 --- a/src/diffusers/pipelines/stable_diffusion_diffedit/pipeline_stable_diffusion_diffedit.py +++ b/src/diffusers/pipelines/stable_diffusion_diffedit/pipeline_stable_diffusion_diffedit.py @@ -249,11 +249,7 @@ class StableDiffusionDiffEditPipeline( StableDiffusionLoraLoaderMixin, ): r""" - - - This is an experimental feature! - - + > [!WARNING] > This is an experimental feature! Pipeline for text-guided image inpainting using Stable Diffusion and DiffEdit. diff --git a/src/diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_k_diffusion.py b/src/diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_k_diffusion.py index df2564a89b..feebd6adf8 100755 --- a/src/diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_k_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_k_diffusion.py @@ -81,11 +81,7 @@ class StableDiffusionKDiffusionPipeline( - [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for loading LoRA weights - [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for saving LoRA weights - - - This is an experimental pipeline and is likely to change in the future. - - + > [!WARNING] > This is an experimental pipeline and is likely to change in the future. Args: vae ([`AutoencoderKL`]): diff --git a/src/diffusers/schedulers/deprecated/scheduling_karras_ve.py b/src/diffusers/schedulers/deprecated/scheduling_karras_ve.py index 6b968e7081..9206ee80a6 100644 --- a/src/diffusers/schedulers/deprecated/scheduling_karras_ve.py +++ b/src/diffusers/schedulers/deprecated/scheduling_karras_ve.py @@ -53,13 +53,9 @@ class KarrasVeScheduler(SchedulerMixin, ConfigMixin): This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic methods the library implements for all schedulers such as loading and saving. - - - For more details on the parameters, see [Appendix E](https://huggingface.co/papers/2206.00364). The grid search - values used to find the optimal `{s_noise, s_churn, s_min, s_max}` for a specific model are described in Table 5 of - the paper. - - + > [!TIP] > For more details on the parameters, see [Appendix E](https://huggingface.co/papers/2206.00364). The grid + search > values used to find the optimal `{s_noise, s_churn, s_min, s_max}` for a specific model are described in + Table 5 of > the paper. Args: sigma_min (`float`, defaults to 0.02): diff --git a/src/diffusers/schedulers/scheduling_consistency_models.py b/src/diffusers/schedulers/scheduling_consistency_models.py index 0f50622588..5d81d5eb8a 100644 --- a/src/diffusers/schedulers/scheduling_consistency_models.py +++ b/src/diffusers/schedulers/scheduling_consistency_models.py @@ -268,11 +268,7 @@ class CMStochasticIterativeScheduler(SchedulerMixin, ConfigMixin): Gets the scalings used in the consistency model parameterization (from Appendix C of the [paper](https://huggingface.co/papers/2303.01469)) to enforce boundary condition. - - - `epsilon` in the equations for `c_skip` and `c_out` is set to `sigma_min`. - - + > [!TIP] > `epsilon` in the equations for `c_skip` and `c_out` is set to `sigma_min`. Args: sigma (`torch.Tensor`): diff --git a/src/diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py b/src/diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py index 66ed296da8..b9567f2c47 100644 --- a/src/diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py +++ b/src/diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py @@ -304,12 +304,8 @@ class CosineDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin): designed to discretize an integral of the noise prediction model, and DPM-Solver++ is designed to discretize an integral of the data prediction model. - - - The algorithm and model type are decoupled. You can use either DPMSolver or DPMSolver++ for both noise - prediction and data prediction models. - - + > [!TIP] > The algorithm and model type are decoupled. You can use either DPMSolver or DPMSolver++ for both + noise > prediction and data prediction models. Args: model_output (`torch.Tensor`): diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py index d07ff8b200..8b523cd13f 100644 --- a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py +++ b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py @@ -630,12 +630,8 @@ class DPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin): designed to discretize an integral of the noise prediction model, and DPM-Solver++ is designed to discretize an integral of the data prediction model. - - - The algorithm and model type are decoupled. You can use either DPMSolver or DPMSolver++ for both noise - prediction and data prediction models. - - + > [!TIP] > The algorithm and model type are decoupled. You can use either DPMSolver or DPMSolver++ for both + noise > prediction and data prediction models. Args: model_output (`torch.Tensor`): diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py b/src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py index 9ec9588511..f1a1ac3d82 100644 --- a/src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py +++ b/src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py @@ -491,12 +491,8 @@ class DPMSolverMultistepInverseScheduler(SchedulerMixin, ConfigMixin): designed to discretize an integral of the noise prediction model, and DPM-Solver++ is designed to discretize an integral of the data prediction model. - - - The algorithm and model type are decoupled. You can use either DPMSolver or DPMSolver++ for both noise - prediction and data prediction models. - - + > [!TIP] > The algorithm and model type are decoupled. You can use either DPMSolver or DPMSolver++ for both + noise > prediction and data prediction models. Args: model_output (`torch.Tensor`): diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py b/src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py index 8663210a62..1ae8249730 100644 --- a/src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py +++ b/src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py @@ -568,12 +568,8 @@ class DPMSolverSinglestepScheduler(SchedulerMixin, ConfigMixin): designed to discretize an integral of the noise prediction model, and DPM-Solver++ is designed to discretize an integral of the data prediction model. - - - The algorithm and model type are decoupled. You can use either DPMSolver or DPMSolver++ for both noise - prediction and data prediction models. - - + > [!TIP] > The algorithm and model type are decoupled. You can use either DPMSolver or DPMSolver++ for both + noise > prediction and data prediction models. Args: model_output (`torch.Tensor`): diff --git a/src/diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py b/src/diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py index f1b38aaff5..e9ba695e1f 100644 --- a/src/diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py +++ b/src/diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py @@ -370,12 +370,8 @@ class EDMDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin): designed to discretize an integral of the noise prediction model, and DPM-Solver++ is designed to discretize an integral of the data prediction model. - - - The algorithm and model type are decoupled. You can use either DPMSolver or DPMSolver++ for both noise - prediction and data prediction models. - - + > [!TIP] > The algorithm and model type are decoupled. You can use either DPMSolver or DPMSolver++ for both + noise > prediction and data prediction models. Args: model_output (`torch.Tensor`): diff --git a/src/diffusers/schedulers/scheduling_sasolver.py b/src/diffusers/schedulers/scheduling_sasolver.py index 2df7d560dd..2979ce193a 100644 --- a/src/diffusers/schedulers/scheduling_sasolver.py +++ b/src/diffusers/schedulers/scheduling_sasolver.py @@ -500,12 +500,8 @@ class SASolverScheduler(SchedulerMixin, ConfigMixin): Noise_prediction is designed to discretize an integral of the noise prediction model, and data_prediction is designed to discretize an integral of the data prediction model. - - - The algorithm and model type are decoupled. You can use either data_prediction or noise_prediction for both - noise prediction and data prediction models. - - + > [!TIP] > The algorithm and model type are decoupled. You can use either data_prediction or noise_prediction + for both > noise prediction and data prediction models. Args: model_output (`torch.Tensor`): diff --git a/src/diffusers/schedulers/scheduling_utils.py b/src/diffusers/schedulers/scheduling_utils.py index f0e162ea6b..a355c7bb1a 100644 --- a/src/diffusers/schedulers/scheduling_utils.py +++ b/src/diffusers/schedulers/scheduling_utils.py @@ -138,15 +138,11 @@ class SchedulerMixin(PushToHubMixin): The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier allowed by Git. - - - To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in with `hf - auth login`. You can also activate the special - ["offline-mode"](https://huggingface.co/diffusers/installation.html#offline-mode) to use this method in a + > [!TIP] > To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in + with `hf > auth login`. You can also activate the special > + ["offline-mode"](https://huggingface.co/diffusers/installation.html#offline-mode) to use this method in a > firewalled environment. - - """ config, kwargs, commit_hash = cls.load_config( pretrained_model_name_or_path=pretrained_model_name_or_path, diff --git a/src/diffusers/schedulers/scheduling_utils_flax.py b/src/diffusers/schedulers/scheduling_utils_flax.py index ffbe3b9020..0534e47d8a 100644 --- a/src/diffusers/schedulers/scheduling_utils_flax.py +++ b/src/diffusers/schedulers/scheduling_utils_flax.py @@ -120,19 +120,12 @@ class FlaxSchedulerMixin(PushToHubMixin): git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any identifier allowed by git. - + > [!TIP] > It is required to be logged in (`hf auth login`) when you want to use private or [gated > + models](https://huggingface.co/docs/hub/models-gated#gated-models). - It is required to be logged in (`hf auth login`) when you want to use private or [gated - models](https://huggingface.co/docs/hub/models-gated#gated-models). - - - - - - Activate the special ["offline-mode"](https://huggingface.co/transformers/installation.html#offline-mode) to - use this method in a firewalled environment. - - + > [!TIP] > Activate the special + ["offline-mode"](https://huggingface.co/transformers/installation.html#offline-mode) to > use this method in a + firewalled environment. """ logger.warning( diff --git a/src/diffusers/utils/dynamic_modules_utils.py b/src/diffusers/utils/dynamic_modules_utils.py index de947a12e2..35bef51229 100644 --- a/src/diffusers/utils/dynamic_modules_utils.py +++ b/src/diffusers/utils/dynamic_modules_utils.py @@ -290,12 +290,8 @@ def get_cached_module_file( local_files_only (`bool`, *optional*, defaults to `False`): If `True`, will only try to load the tokenizer configuration from local files. - - - You may pass a token in `token` if you are not logged in (`hf auth login`) and want to use private or [gated - models](https://huggingface.co/docs/hub/models-gated#gated-models). - - + > [!TIP] > You may pass a token in `token` if you are not logged in (`hf auth login`) and want to use private or + [gated > models](https://huggingface.co/docs/hub/models-gated#gated-models). Returns: `str`: The path to the module inside the cache. @@ -440,12 +436,8 @@ def get_class_from_dynamic_module( """ Extracts a class from a module file, present in the local folder or repository of a model. - - - Calling this function will execute the code in the module file found locally or downloaded from the Hub. It should - therefore only be called on trusted repos. - - + > [!WARNING] > Calling this function will execute the code in the module file found locally or downloaded from the + Hub. It should > therefore only be called on trusted repos. Args: pretrained_model_name_or_path (`str` or `os.PathLike`): @@ -480,12 +472,8 @@ def get_class_from_dynamic_module( local_files_only (`bool`, *optional*, defaults to `False`): If `True`, will only try to load the tokenizer configuration from local files. - - - You may pass a token in `token` if you are not logged in (`hf auth login`) and want to use private or [gated - models](https://huggingface.co/docs/hub/models-gated#gated-models). - - + > [!TIP] > You may pass a token in `token` if you are not logged in (`hf auth login`) and want to use private or + [gated > models](https://huggingface.co/docs/hub/models-gated#gated-models). Returns: `type`: The class, dynamically imported from the module. diff --git a/src/diffusers/utils/outputs.py b/src/diffusers/utils/outputs.py index 35691496a1..2b20f6120c 100644 --- a/src/diffusers/utils/outputs.py +++ b/src/diffusers/utils/outputs.py @@ -43,12 +43,8 @@ class BaseOutput(OrderedDict): tuple) or strings (like a dictionary) that will ignore the `None` attributes. Otherwise behaves like a regular Python dictionary. - - - You can't unpack a [`BaseOutput`] directly. Use the [`~utils.BaseOutput.to_tuple`] method to convert it to a tuple - first. - - + > [!WARNING] > You can't unpack a [`BaseOutput`] directly. Use the [`~utils.BaseOutput.to_tuple`] method to convert + it to a tuple > first. """ def __init_subclass__(cls) -> None: From 814d710e56fb072542120662f0d22e2556168a91 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Wed, 1 Oct 2025 09:02:55 +0530 Subject: [PATCH 042/123] [tests] cache non lora pipeline outputs. (#12298) * cache non lora pipeline outputs. * up * up * up * up * Revert "up" This reverts commit 772c32e43397f25919c29bbbe8ef9dc7d581cfb8. * up * Revert "up" This reverts commit cca03df7fce55550ed28b59cadec12d1db188283. * up * up * add . * up * up * up * up * up * up --- tests/lora/test_lora_layers_cogview4.py | 3 - tests/lora/test_lora_layers_flux.py | 11 +-- tests/lora/test_lora_layers_wanvace.py | 2 +- tests/lora/utils.py | 109 +++++++++--------------- 4 files changed, 42 insertions(+), 83 deletions(-) diff --git a/tests/lora/test_lora_layers_cogview4.py b/tests/lora/test_lora_layers_cogview4.py index 9c62d2f0b8..30eb8fbb63 100644 --- a/tests/lora/test_lora_layers_cogview4.py +++ b/tests/lora/test_lora_layers_cogview4.py @@ -129,9 +129,6 @@ class CogView4LoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) - images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] with tempfile.TemporaryDirectory() as tmpdirname: diff --git a/tests/lora/test_lora_layers_flux.py b/tests/lora/test_lora_layers_flux.py index 6c22a34889..b840d7ac72 100644 --- a/tests/lora/test_lora_layers_flux.py +++ b/tests/lora/test_lora_layers_flux.py @@ -122,9 +122,6 @@ class FluxLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0)).images - self.assertTrue(output_no_lora.shape == self.output_shape) - pipe.transformer.add_adapter(denoiser_lora_config) self.assertTrue(check_if_lora_correctly_set(pipe.transformer), "Lora not correctly set in transformer") @@ -170,8 +167,7 @@ class FluxLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0)).images - self.assertTrue(output_no_lora.shape == self.output_shape) + output_no_lora = self.get_base_pipe_output() # Modify the config to have a layer which won't be present in the second LoRA we will load. modified_denoiser_lora_config = copy.deepcopy(denoiser_lora_config) @@ -218,9 +214,7 @@ class FluxLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): pipe = pipe.to(torch_device) pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0)).images - self.assertTrue(output_no_lora.shape == self.output_shape) + output_no_lora = self.get_base_pipe_output() # Modify the config to have a layer which won't be present in the first LoRA we will load. modified_denoiser_lora_config = copy.deepcopy(denoiser_lora_config) @@ -329,6 +323,7 @@ class FluxControlLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): noise = floats_tensor((batch_size, num_channels) + sizes) input_ids = torch.randint(1, sequence_length, size=(batch_size, sequence_length), generator=generator) + np.random.seed(0) pipeline_inputs = { "prompt": "A painting of a squirrel eating a burger", "control_image": Image.fromarray(np.random.randint(0, 255, size=(32, 32, 3), dtype="uint8")), diff --git a/tests/lora/test_lora_layers_wanvace.py b/tests/lora/test_lora_layers_wanvace.py index c3244e150e..ab1f57bfc9 100644 --- a/tests/lora/test_lora_layers_wanvace.py +++ b/tests/lora/test_lora_layers_wanvace.py @@ -169,7 +169,7 @@ class WanVACELoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): pipe = self.pipeline_class(**components).to(torch_device) _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + output_no_lora = self.get_base_pipe_output() self.assertTrue(output_no_lora.shape == self.output_shape) # only supported for `denoiser` now diff --git a/tests/lora/utils.py b/tests/lora/utils.py index ecaa553ce4..3d4344bb86 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -126,13 +126,20 @@ class PeftLoraLoaderMixinTests: text_encoder_target_modules = ["q_proj", "k_proj", "v_proj", "out_proj"] denoiser_target_modules = ["to_q", "to_k", "to_v", "to_out.0"] - def get_dummy_components(self, use_dora=False, lora_alpha=None): + cached_non_lora_output = None + + def get_base_pipe_output(self): + if self.cached_non_lora_output is None: + self.cached_non_lora_output = self._compute_baseline_output() + return self.cached_non_lora_output + + def get_dummy_components(self, scheduler_cls=None, use_dora=False, lora_alpha=None): if self.unet_kwargs and self.transformer_kwargs: raise ValueError("Both `unet_kwargs` and `transformer_kwargs` cannot be specified.") if self.has_two_text_encoders and self.has_three_text_encoders: raise ValueError("Both `has_two_text_encoders` and `has_three_text_encoders` cannot be True.") - scheduler_cls = self.scheduler_cls + scheduler_cls = scheduler_cls if scheduler_cls is not None else self.scheduler_cls rank = 4 lora_alpha = rank if lora_alpha is None else lora_alpha @@ -238,15 +245,16 @@ class PeftLoraLoaderMixinTests: return noise, input_ids, pipeline_inputs - # Copied from: https://colab.research.google.com/gist/sayakpaul/df2ef6e1ae6d8c10a49d859883b10860/scratchpad.ipynb - def get_dummy_tokens(self): - max_seq_length = 77 + def _compute_baseline_output(self): + components, _, _ = self.get_dummy_components(self.scheduler_cls) + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) - inputs = torch.randint(2, 56, size=(1, max_seq_length), generator=torch.manual_seed(0)) - - prepared_inputs = {} - prepared_inputs["input_ids"] = inputs - return prepared_inputs + # Always ensure the inputs are without the `generator`. Make sure to pass the `generator` + # explicitly. + _, _, inputs = self.get_dummy_inputs(with_generator=False) + return pipe(**inputs, generator=torch.manual_seed(0))[0] def _get_lora_state_dicts(self, modules_to_save): state_dicts = {} @@ -316,14 +324,8 @@ class PeftLoraLoaderMixinTests: """ Tests a simple inference and makes sure it works as expected """ - components, text_lora_config, _ = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - _, _, inputs = self.get_dummy_inputs() - output_no_lora = pipe(**inputs)[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + output_no_lora = self.get_base_pipe_output() + assert output_no_lora.shape == self.output_shape def test_simple_inference_with_text_lora(self): """ @@ -336,9 +338,7 @@ class PeftLoraLoaderMixinTests: pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) - + output_no_lora = self.get_base_pipe_output() pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] @@ -414,9 +414,6 @@ class PeftLoraLoaderMixinTests: pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] @@ -466,8 +463,7 @@ class PeftLoraLoaderMixinTests: pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + output_no_lora = self.get_base_pipe_output() pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) @@ -503,8 +499,7 @@ class PeftLoraLoaderMixinTests: pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + output_no_lora = self.get_base_pipe_output() pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) @@ -534,8 +529,7 @@ class PeftLoraLoaderMixinTests: pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + output_no_lora = self.get_base_pipe_output() pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) @@ -566,9 +560,6 @@ class PeftLoraLoaderMixinTests: pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] @@ -616,8 +607,7 @@ class PeftLoraLoaderMixinTests: pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + output_no_lora = self.get_base_pipe_output() pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) @@ -666,9 +656,6 @@ class PeftLoraLoaderMixinTests: pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] @@ -708,9 +695,6 @@ class PeftLoraLoaderMixinTests: pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] @@ -747,9 +731,7 @@ class PeftLoraLoaderMixinTests: pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) - + output_no_lora = self.get_base_pipe_output() pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] @@ -790,8 +772,7 @@ class PeftLoraLoaderMixinTests: pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + output_no_lora = self.get_base_pipe_output() pipe, denoiser = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) @@ -825,8 +806,7 @@ class PeftLoraLoaderMixinTests: pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + output_no_lora = self.get_base_pipe_output() pipe, denoiser = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) @@ -900,7 +880,7 @@ class PeftLoraLoaderMixinTests: pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + output_no_lora = self.get_base_pipe_output() if "text_encoder" in self.pipeline_class._lora_loadable_modules: pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") @@ -1024,7 +1004,7 @@ class PeftLoraLoaderMixinTests: pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + output_no_lora = self.get_base_pipe_output() pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") @@ -1080,7 +1060,7 @@ class PeftLoraLoaderMixinTests: pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + output_no_lora = self.get_base_pipe_output() if "text_encoder" in self.pipeline_class._lora_loadable_modules: pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") @@ -1240,7 +1220,7 @@ class PeftLoraLoaderMixinTests: pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + output_no_lora = self.get_base_pipe_output() if "text_encoder" in self.pipeline_class._lora_loadable_modules: pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") @@ -1331,7 +1311,7 @@ class PeftLoraLoaderMixinTests: pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + output_no_lora = self.get_base_pipe_output() if "text_encoder" in self.pipeline_class._lora_loadable_modules: pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") @@ -1551,7 +1531,6 @@ class PeftLoraLoaderMixinTests: self.assertDictEqual(pipe.get_list_adapters(), dicts_to_be_checked) - @require_peft_version_greater(peft_version="0.6.2") def test_simple_inference_with_text_lora_denoiser_fused_multi( self, expected_atol: float = 1e-3, expected_rtol: float = 1e-3 ): @@ -1565,9 +1544,6 @@ class PeftLoraLoaderMixinTests: pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) - if "text_encoder" in self.pipeline_class._lora_loadable_modules: pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") @@ -1641,8 +1617,7 @@ class PeftLoraLoaderMixinTests: pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + output_no_lora = self.get_base_pipe_output() if "text_encoder" in self.pipeline_class._lora_loadable_modules: pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") @@ -1685,7 +1660,6 @@ class PeftLoraLoaderMixinTests: "LoRA should change the output", ) - @require_peft_version_greater(peft_version="0.9.0") def test_simple_inference_with_dora(self): components, text_lora_config, denoiser_lora_config = self.get_dummy_components(use_dora=True) pipe = self.pipeline_class(**components) @@ -1695,7 +1669,6 @@ class PeftLoraLoaderMixinTests: output_no_dora_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] self.assertTrue(output_no_dora_lora.shape == self.output_shape) - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) output_dora_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] @@ -1783,7 +1756,6 @@ class PeftLoraLoaderMixinTests: pipe = pipe.to(torch_device) pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True) @@ -1820,7 +1792,7 @@ class PeftLoraLoaderMixinTests: pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - original_out = pipe(**inputs, generator=torch.manual_seed(0))[0] + output_no_lora = self.get_base_pipe_output() no_op_state_dict = {"lora_foo": torch.tensor(2.0), "lora_bar": torch.tensor(3.0)} logger = logging.get_logger("diffusers.loaders.peft") @@ -1832,7 +1804,7 @@ class PeftLoraLoaderMixinTests: denoiser = getattr(pipe, "unet") if self.unet_kwargs is not None else getattr(pipe, "transformer") self.assertTrue(cap_logger.out.startswith(f"No LoRA keys associated to {denoiser.__class__.__name__}")) - self.assertTrue(np.allclose(original_out, out_after_lora_attempt, atol=1e-5, rtol=1e-5)) + self.assertTrue(np.allclose(output_no_lora, out_after_lora_attempt, atol=1e-5, rtol=1e-5)) # test only for text encoder for lora_module in self.pipeline_class._lora_loadable_modules: @@ -1864,9 +1836,7 @@ class PeftLoraLoaderMixinTests: pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) - + output_no_lora = self.get_base_pipe_output() pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) lora_scale = 0.5 @@ -2212,9 +2182,6 @@ class PeftLoraLoaderMixinTests: pipe = self.pipeline_class(**components).to(torch_device) _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) - pipe, _ = self.add_adapters_to_pipeline( pipe, text_lora_config=text_lora_config, denoiser_lora_config=denoiser_lora_config ) @@ -2260,7 +2227,7 @@ class PeftLoraLoaderMixinTests: pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + output_no_lora = self.get_base_pipe_output() if "text_encoder" in self.pipeline_class._lora_loadable_modules: pipe.text_encoder.add_adapter(text_lora_config) From 9ae5b6299d3b1a7b0378dc77c3c69baf521587d2 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 2 Oct 2025 17:46:15 +0530 Subject: [PATCH 043/123] [ci] xfail failing tests in CI. (#12418) xfail failing tests in CI. --- tests/pipelines/test_pipelines.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/pipelines/test_pipelines.py b/tests/pipelines/test_pipelines.py index 3a69813612..5b86423553 100644 --- a/tests/pipelines/test_pipelines.py +++ b/tests/pipelines/test_pipelines.py @@ -28,6 +28,7 @@ import warnings import numpy as np import PIL.Image +import pytest import requests_mock import safetensors.torch import torch @@ -62,10 +63,7 @@ from diffusers import ( ) from diffusers.pipelines.pipeline_utils import _get_pipeline_class from diffusers.schedulers.scheduling_utils import SCHEDULER_CONFIG_NAME -from diffusers.utils import ( - CONFIG_NAME, - WEIGHTS_NAME, -) +from diffusers.utils import CONFIG_NAME, WEIGHTS_NAME, is_transformers_version from diffusers.utils.torch_utils import is_compiled_module from ..testing_utils import ( @@ -584,6 +582,7 @@ class DownloadTests(unittest.TestCase): assert not any(f.endswith(unexpected_ext) for f in files) assert all(variant in f for f in model_files if f.endswith(model_ext) and variant is not None) + @pytest.mark.xfail(condition=is_transformers_version(">", "4.56.2"), reason="Some import error", strict=True) def test_download_legacy_variants_with_sharded_ckpts_raises_warning(self): repo_id = "hf-internal-testing/tiny-stable-diffusion-pipe-variants-all-kinds" logger = logging.get_logger("diffusers.pipelines.pipeline_utils") @@ -690,6 +689,7 @@ class DownloadTests(unittest.TestCase): ) assert "Error no file name" in str(error_context.exception) + @pytest.mark.xfail(condition=is_transformers_version(">", "4.56.2"), reason="Some import error", strict=True) def test_local_save_load_index(self): prompt = "hello" for variant in [None, "fp16"]: @@ -1584,6 +1584,7 @@ class PipelineFastTests(unittest.TestCase): assert pipeline.scheduler is not None assert pipeline.feature_extractor is not None + @pytest.mark.xfail(condition=is_transformers_version(">", "4.56.2"), reason="Some import error", strict=True) def test_no_pytorch_download_when_doing_safetensors(self): # by default we don't download with tempfile.TemporaryDirectory() as tmpdirname: @@ -1603,6 +1604,7 @@ class PipelineFastTests(unittest.TestCase): # pytorch does not assert not os.path.exists(os.path.join(path, "diffusion_pytorch_model.bin")) + @pytest.mark.xfail(condition=is_transformers_version(">", "4.56.2"), reason="Some import error", strict=True) def test_no_safetensors_download_when_doing_pytorch(self): use_safetensors = False From b4297967a04cca6ac4493202c02d81c30d0f9ee8 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 2 Oct 2025 20:38:02 +0530 Subject: [PATCH 044/123] [core] conditionally import torch distributed stuff. (#12420) conditionally import torch distributed stuff. --- src/diffusers/models/attention_dispatch.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index 0a2ad68123..e169491099 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -20,7 +20,10 @@ from enum import Enum from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Tuple, Union import torch -import torch.distributed._functional_collectives as funcol + + +if torch.distributed.is_available(): + import torch.distributed._functional_collectives as funcol from ..utils import ( get_logger, From 7242b5ff627fad93dd85834b0278267b6cbe2d6d Mon Sep 17 00:00:00 2001 From: Benjamin Bossan Date: Thu, 2 Oct 2025 20:57:11 +0200 Subject: [PATCH 045/123] FIX Test to ignore warning for enable_lora_hotswap (#12421) I noticed that the test should be for the option check_compiled="ignore" but it was using check_compiled="warn". This has been fixed, now the correct argument is passed. However, the fact that the test passed means that it was incorrect to begin with. The way that logs are collected does not collect the logger.warning call here (not sure why). To amend this, I'm now using assertNoLogs. With this change, the test correctly fails when the wrong argument is passed. --- tests/models/test_modeling_common.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/models/test_modeling_common.py b/tests/models/test_modeling_common.py index 9b1c6b50dc..a44ef571c5 100644 --- a/tests/models/test_modeling_common.py +++ b/tests/models/test_modeling_common.py @@ -25,7 +25,6 @@ import traceback import unittest import unittest.mock as mock import uuid -import warnings from collections import defaultdict from typing import Dict, List, Optional, Tuple, Union @@ -2373,14 +2372,15 @@ class LoraHotSwappingForModelTesterMixin: def test_enable_lora_hotswap_called_after_adapter_added_ignore(self): # check possibility to ignore the error/warning + from diffusers.loaders.peft import logger + lora_config = self.get_lora_config(8, 8, target_modules=["to_q"]) init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() model = self.model_class(**init_dict).to(torch_device) model.add_adapter(lora_config) - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") # Capture all warnings - model.enable_lora_hotswap(target_rank=32, check_compiled="warn") - self.assertEqual(len(w), 0, f"Expected no warnings, but got: {[str(warn.message) for warn in w]}") + # note: assertNoLogs requires Python 3.10+ + with self.assertNoLogs(logger, level="WARNING"): + model.enable_lora_hotswap(target_rank=32, check_compiled="ignore") def test_enable_lora_hotswap_wrong_check_compiled_argument_raises(self): # check that wrong argument value raises an error From 941ac9c3d9aab9c36fc33c58dac1980442928082 Mon Sep 17 00:00:00 2001 From: Linoy Tsaban <57615435+linoytsaban@users.noreply.github.com> Date: Fri, 3 Oct 2025 17:46:47 +0300 Subject: [PATCH 046/123] [training-scripts] Make more examples UV-compatible (follow up on #12000) (#12407) * make qwen and kontext uv compatible * add torchvision * add torchvision * add datasets, bitsandbytes, prodigyopt --------- Co-authored-by: Sayak Paul --- .../train_dreambooth_lora_flux_advanced.py | 4 ++++ .../dreambooth/train_dreambooth_lora_flux.py | 4 ++++ .../train_dreambooth_lora_flux_kontext.py | 18 ++++++++++++++++++ .../train_dreambooth_lora_qwen_image.py | 18 ++++++++++++++++++ .../dreambooth/train_dreambooth_lora_sana.py | 4 ++++ 5 files changed, 48 insertions(+) diff --git a/examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced.py b/examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced.py index a46490e8b3..5aa33190d4 100644 --- a/examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced.py +++ b/examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced.py @@ -25,6 +25,10 @@ # "Jinja2", # "peft>=0.11.1", # "sentencepiece", +# "torchvision", +# "datasets", +# "bitsandbytes", +# "prodigyopt", # ] # /// diff --git a/examples/dreambooth/train_dreambooth_lora_flux.py b/examples/dreambooth/train_dreambooth_lora_flux.py index bd3a974a17..3b6ab814f2 100644 --- a/examples/dreambooth/train_dreambooth_lora_flux.py +++ b/examples/dreambooth/train_dreambooth_lora_flux.py @@ -25,6 +25,10 @@ # "Jinja2", # "peft>=0.11.1", # "sentencepiece", +# "torchvision", +# "datasets", +# "bitsandbytes", +# "prodigyopt", # ] # /// diff --git a/examples/dreambooth/train_dreambooth_lora_flux_kontext.py b/examples/dreambooth/train_dreambooth_lora_flux_kontext.py index 03c05a05e0..fc6df87768 100644 --- a/examples/dreambooth/train_dreambooth_lora_flux_kontext.py +++ b/examples/dreambooth/train_dreambooth_lora_flux_kontext.py @@ -14,6 +14,24 @@ # See the License for the specific language governing permissions and # limitations under the License. +# /// script +# dependencies = [ +# "diffusers @ git+https://github.com/huggingface/diffusers.git", +# "torch>=2.0.0", +# "accelerate>=0.31.0", +# "transformers>=4.41.2", +# "ftfy", +# "tensorboard", +# "Jinja2", +# "peft>=0.11.1", +# "sentencepiece", +# "torchvision", +# "datasets", +# "bitsandbytes", +# "prodigyopt", +# ] +# /// + import argparse import copy import itertools diff --git a/examples/dreambooth/train_dreambooth_lora_qwen_image.py b/examples/dreambooth/train_dreambooth_lora_qwen_image.py index feec4da712..75eae92dfb 100644 --- a/examples/dreambooth/train_dreambooth_lora_qwen_image.py +++ b/examples/dreambooth/train_dreambooth_lora_qwen_image.py @@ -13,6 +13,24 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and +# /// script +# dependencies = [ +# "diffusers @ git+https://github.com/huggingface/diffusers.git", +# "torch>=2.0.0", +# "accelerate>=0.31.0", +# "transformers>=4.41.2", +# "ftfy", +# "tensorboard", +# "Jinja2", +# "peft>=0.11.1", +# "sentencepiece", +# "torchvision", +# "datasets", +# "bitsandbytes", +# "prodigyopt", +# ] +# /// + import argparse import copy import itertools diff --git a/examples/dreambooth/train_dreambooth_lora_sana.py b/examples/dreambooth/train_dreambooth_lora_sana.py index b188a80916..2b0c1ee669 100644 --- a/examples/dreambooth/train_dreambooth_lora_sana.py +++ b/examples/dreambooth/train_dreambooth_lora_sana.py @@ -25,6 +25,10 @@ # "Jinja2", # "peft>=0.14.0", # "sentencepiece", +# "torchvision", +# "datasets", +# "bitsandbytes", +# "prodigyopt", # ] # /// From 2b7deffe361b7b0e1d2665a1f9f0bd4daea4927c Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sat, 4 Oct 2025 23:53:38 -0400 Subject: [PATCH 047/123] fix scale_shift_factor being on cpu for wan and ltx (#12347) * wan fix scale_shift_factor being on cpu * apply device cast to ltx transformer * Apply style fixes --------- Co-authored-by: Dhruv Nair Co-authored-by: github-actions[bot] --- src/diffusers/models/transformers/transformer_ltx.py | 4 +++- src/diffusers/models/transformers/transformer_wan.py | 4 ++-- src/diffusers/models/transformers/transformer_wan_vace.py | 4 ++-- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_ltx.py b/src/diffusers/models/transformers/transformer_ltx.py index 9f3840690d..685c73c07c 100644 --- a/src/diffusers/models/transformers/transformer_ltx.py +++ b/src/diffusers/models/transformers/transformer_ltx.py @@ -353,7 +353,9 @@ class LTXVideoTransformerBlock(nn.Module): norm_hidden_states = self.norm1(hidden_states) num_ada_params = self.scale_shift_table.shape[0] - ada_values = self.scale_shift_table[None, None] + temb.reshape(batch_size, temb.size(1), num_ada_params, -1) + ada_values = self.scale_shift_table[None, None].to(temb.device) + temb.reshape( + batch_size, temb.size(1), num_ada_params, -1 + ) shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ada_values.unbind(dim=2) norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa diff --git a/src/diffusers/models/transformers/transformer_wan.py b/src/diffusers/models/transformers/transformer_wan.py index 25c055fb56..dd75fb124f 100644 --- a/src/diffusers/models/transformers/transformer_wan.py +++ b/src/diffusers/models/transformers/transformer_wan.py @@ -682,12 +682,12 @@ class WanTransformer3DModel( # 5. Output norm, projection & unpatchify if temb.ndim == 3: # batch_size, seq_len, inner_dim (wan 2.2 ti2v) - shift, scale = (self.scale_shift_table.unsqueeze(0) + temb.unsqueeze(2)).chunk(2, dim=2) + shift, scale = (self.scale_shift_table.unsqueeze(0).to(temb.device) + temb.unsqueeze(2)).chunk(2, dim=2) shift = shift.squeeze(2) scale = scale.squeeze(2) else: # batch_size, inner_dim - shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1) + shift, scale = (self.scale_shift_table.to(temb.device) + temb.unsqueeze(1)).chunk(2, dim=1) # Move the shift and scale tensors to the same device as hidden_states. # When using multi-GPU inference via accelerate these will be on the diff --git a/src/diffusers/models/transformers/transformer_wan_vace.py b/src/diffusers/models/transformers/transformer_wan_vace.py index e5a9c7e0a6..30c38c244a 100644 --- a/src/diffusers/models/transformers/transformer_wan_vace.py +++ b/src/diffusers/models/transformers/transformer_wan_vace.py @@ -103,7 +103,7 @@ class WanVACETransformerBlock(nn.Module): control_hidden_states = control_hidden_states + hidden_states shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( - self.scale_shift_table + temb.float() + self.scale_shift_table.to(temb.device) + temb.float() ).chunk(6, dim=1) # 1. Self-attention @@ -361,7 +361,7 @@ class WanVACETransformer3DModel( hidden_states = hidden_states + control_hint * scale # 6. Output norm, projection & unpatchify - shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1) + shift, scale = (self.scale_shift_table.to(temb.device) + temb.unsqueeze(1)).chunk(2, dim=1) # Move the shift and scale tensors to the same device as hidden_states. # When using multi-GPU inference via accelerate these will be on the From c3675d4c9bb9c02521cd2c1aec198460c1657256 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Sun, 5 Oct 2025 21:57:13 +0530 Subject: [PATCH 048/123] [core] support QwenImage Edit Plus in modular (#12416) * up * up * up * up * up * up * remove saves * move things around a bit. * get ready. --- src/diffusers/__init__.py | 4 + src/diffusers/modular_pipelines/__init__.py | 4 + .../modular_pipelines/modular_pipeline.py | 4 +- .../modular_pipelines/qwenimage/__init__.py | 18 +- .../qwenimage/before_denoise.py | 3 +- .../modular_pipelines/qwenimage/encoders.py | 236 +++++++++++++++++- .../qwenimage/modular_blocks.py | 151 ++++++++++- .../qwenimage/modular_pipeline.py | 10 + src/diffusers/pipelines/auto_pipeline.py | 2 + .../dummy_torch_and_transformers_objects.py | 30 +++ 10 files changed, 449 insertions(+), 13 deletions(-) diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 8867250ded..686e8d99da 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -390,6 +390,8 @@ else: "QwenImageAutoBlocks", "QwenImageEditAutoBlocks", "QwenImageEditModularPipeline", + "QwenImageEditPlusAutoBlocks", + "QwenImageEditPlusModularPipeline", "QwenImageModularPipeline", "StableDiffusionXLAutoBlocks", "StableDiffusionXLModularPipeline", @@ -1052,6 +1054,8 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: QwenImageAutoBlocks, QwenImageEditAutoBlocks, QwenImageEditModularPipeline, + QwenImageEditPlusAutoBlocks, + QwenImageEditPlusModularPipeline, QwenImageModularPipeline, StableDiffusionXLAutoBlocks, StableDiffusionXLModularPipeline, diff --git a/src/diffusers/modular_pipelines/__init__.py b/src/diffusers/modular_pipelines/__init__.py index 65c22b349b..2e590594af 100644 --- a/src/diffusers/modular_pipelines/__init__.py +++ b/src/diffusers/modular_pipelines/__init__.py @@ -52,6 +52,8 @@ else: "QwenImageModularPipeline", "QwenImageEditModularPipeline", "QwenImageEditAutoBlocks", + "QwenImageEditPlusModularPipeline", + "QwenImageEditPlusAutoBlocks", ] _import_structure["components_manager"] = ["ComponentsManager"] @@ -78,6 +80,8 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: QwenImageAutoBlocks, QwenImageEditAutoBlocks, QwenImageEditModularPipeline, + QwenImageEditPlusAutoBlocks, + QwenImageEditPlusModularPipeline, QwenImageModularPipeline, ) from .stable_diffusion_xl import StableDiffusionXLAutoBlocks, StableDiffusionXLModularPipeline diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 037c9e323c..e543bf0bb3 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -59,6 +59,7 @@ MODULAR_PIPELINE_MAPPING = OrderedDict( ("flux", "FluxModularPipeline"), ("qwenimage", "QwenImageModularPipeline"), ("qwenimage-edit", "QwenImageEditModularPipeline"), + ("qwenimage-edit-plus", "QwenImageEditPlusModularPipeline"), ] ) @@ -1628,7 +1629,8 @@ class ModularPipeline(ConfigMixin, PushToHubMixin): blocks = ModularPipelineBlocks.from_pretrained( pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **kwargs ) - except EnvironmentError: + except EnvironmentError as e: + logger.debug(f"EnvironmentError: {e}") blocks = None cache_dir = kwargs.pop("cache_dir", None) diff --git a/src/diffusers/modular_pipelines/qwenimage/__init__.py b/src/diffusers/modular_pipelines/qwenimage/__init__.py index 81cf515730..ae4ec4799f 100644 --- a/src/diffusers/modular_pipelines/qwenimage/__init__.py +++ b/src/diffusers/modular_pipelines/qwenimage/__init__.py @@ -29,13 +29,20 @@ else: "EDIT_AUTO_BLOCKS", "EDIT_BLOCKS", "EDIT_INPAINT_BLOCKS", + "EDIT_PLUS_AUTO_BLOCKS", + "EDIT_PLUS_BLOCKS", "IMAGE2IMAGE_BLOCKS", "INPAINT_BLOCKS", "TEXT2IMAGE_BLOCKS", "QwenImageAutoBlocks", "QwenImageEditAutoBlocks", + "QwenImageEditPlusAutoBlocks", + ] + _import_structure["modular_pipeline"] = [ + "QwenImageEditModularPipeline", + "QwenImageEditPlusModularPipeline", + "QwenImageModularPipeline", ] - _import_structure["modular_pipeline"] = ["QwenImageEditModularPipeline", "QwenImageModularPipeline"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: try: @@ -54,13 +61,20 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: EDIT_AUTO_BLOCKS, EDIT_BLOCKS, EDIT_INPAINT_BLOCKS, + EDIT_PLUS_AUTO_BLOCKS, + EDIT_PLUS_BLOCKS, IMAGE2IMAGE_BLOCKS, INPAINT_BLOCKS, TEXT2IMAGE_BLOCKS, QwenImageAutoBlocks, QwenImageEditAutoBlocks, + QwenImageEditPlusAutoBlocks, + ) + from .modular_pipeline import ( + QwenImageEditModularPipeline, + QwenImageEditPlusModularPipeline, + QwenImageModularPipeline, ) - from .modular_pipeline import QwenImageEditModularPipeline, QwenImageModularPipeline else: import sys diff --git a/src/diffusers/modular_pipelines/qwenimage/before_denoise.py b/src/diffusers/modular_pipelines/qwenimage/before_denoise.py index 606236cfe9..fdec95dc50 100644 --- a/src/diffusers/modular_pipelines/qwenimage/before_denoise.py +++ b/src/diffusers/modular_pipelines/qwenimage/before_denoise.py @@ -203,7 +203,6 @@ class QwenImagePrepareLatentsStep(ModularPipelineBlocks): block_state.latents = components.pachifier.pack_latents(block_state.latents) self.set_block_state(state, block_state) - return components, state @@ -571,7 +570,7 @@ class QwenImageEditRoPEInputsStep(ModularPipelineBlocks): @property def description(self) -> str: - return "Step that prepares the RoPE inputs for denoising process. This is used in QwenImage Edit. Should be place after prepare_latents step" + return "Step that prepares the RoPE inputs for denoising process. This is used in QwenImage Edit. Should be placed after prepare_latents step" @property def inputs(self) -> List[InputParam]: diff --git a/src/diffusers/modular_pipelines/qwenimage/encoders.py b/src/diffusers/modular_pipelines/qwenimage/encoders.py index 2ab83a03ee..04fb3fdc94 100644 --- a/src/diffusers/modular_pipelines/qwenimage/encoders.py +++ b/src/diffusers/modular_pipelines/qwenimage/encoders.py @@ -128,6 +128,61 @@ def get_qwen_prompt_embeds_edit( return prompt_embeds, encoder_attention_mask +def get_qwen_prompt_embeds_edit_plus( + text_encoder, + processor, + prompt: Union[str, List[str]] = None, + image: Optional[Union[torch.Tensor, List[PIL.Image.Image], PIL.Image.Image]] = None, + prompt_template_encode: str = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n", + img_template_encode: str = "Picture {}: <|vision_start|><|image_pad|><|vision_end|>", + prompt_template_encode_start_idx: int = 64, + device: Optional[torch.device] = None, +): + prompt = [prompt] if isinstance(prompt, str) else prompt + if isinstance(image, list): + base_img_prompt = "" + for i, img in enumerate(image): + base_img_prompt += img_template_encode.format(i + 1) + elif image is not None: + base_img_prompt = img_template_encode.format(1) + else: + base_img_prompt = "" + + template = prompt_template_encode + + drop_idx = prompt_template_encode_start_idx + txt = [template.format(base_img_prompt + e) for e in prompt] + + model_inputs = processor( + text=txt, + images=image, + padding=True, + return_tensors="pt", + ).to(device) + outputs = text_encoder( + input_ids=model_inputs.input_ids, + attention_mask=model_inputs.attention_mask, + pixel_values=model_inputs.pixel_values, + image_grid_thw=model_inputs.image_grid_thw, + output_hidden_states=True, + ) + + hidden_states = outputs.hidden_states[-1] + split_hidden_states = _extract_masked_hidden(hidden_states, model_inputs.attention_mask) + split_hidden_states = [e[drop_idx:] for e in split_hidden_states] + attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states] + max_seq_len = max([e.size(0) for e in split_hidden_states]) + prompt_embeds = torch.stack( + [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states] + ) + encoder_attention_mask = torch.stack( + [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list] + ) + + prompt_embeds = prompt_embeds.to(device=device) + return prompt_embeds, encoder_attention_mask + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" @@ -266,6 +321,83 @@ class QwenImageEditResizeDynamicStep(ModularPipelineBlocks): return components, state +class QwenImageEditPlusResizeDynamicStep(QwenImageEditResizeDynamicStep): + model_name = "qwenimage" + + def __init__( + self, + input_name: str = "image", + output_name: str = "resized_image", + vae_image_output_name: str = "vae_image", + ): + """Create a configurable step for resizing images to the target area (1024 * 1024) while maintaining the aspect ratio. + + This block resizes an input image or a list input images and exposes the resized result under configurable + input and output names. Use this when you need to wire the resize step to different image fields (e.g., + "image", "control_image") + + Args: + input_name (str, optional): Name of the image field to read from the + pipeline state. Defaults to "image". + output_name (str, optional): Name of the resized image field to write + back to the pipeline state. Defaults to "resized_image". + vae_image_output_name (str, optional): Name of the image field + to write back to the pipeline state. This is used by the VAE encoder step later on. QwenImage Edit Plus + processes the input image(s) differently for the VL and the VAE. + """ + if not isinstance(input_name, str) or not isinstance(output_name, str): + raise ValueError( + f"input_name and output_name must be strings but are {type(input_name)} and {type(output_name)}" + ) + self.condition_image_size = 384 * 384 + self._image_input_name = input_name + self._resized_image_output_name = output_name + self._vae_image_output_name = vae_image_output_name + super().__init__() + + @property + def intermediate_outputs(self) -> List[OutputParam]: + return super().intermediate_outputs + [ + OutputParam( + name=self._vae_image_output_name, + type_hint=List[PIL.Image.Image], + description="The images to be processed which will be further used by the VAE encoder.", + ), + ] + + @torch.no_grad() + def __call__(self, components: QwenImageModularPipeline, state: PipelineState): + block_state = self.get_block_state(state) + + images = getattr(block_state, self._image_input_name) + + if not is_valid_image_imagelist(images): + raise ValueError(f"Images must be image or list of images but are {type(images)}") + + if ( + not isinstance(images, torch.Tensor) + and isinstance(images, PIL.Image.Image) + and not isinstance(images, list) + ): + images = [images] + + # TODO (sayakpaul): revisit this when the inputs are `torch.Tensor`s + condition_images = [] + vae_images = [] + for img in images: + image_width, image_height = img.size + condition_width, condition_height, _ = calculate_dimensions( + self.condition_image_size, image_width / image_height + ) + condition_images.append(components.image_resize_processor.resize(img, condition_height, condition_width)) + vae_images.append(img) + + setattr(block_state, self._resized_image_output_name, condition_images) + setattr(block_state, self._vae_image_output_name, vae_images) + self.set_block_state(state, block_state) + return components, state + + class QwenImageTextEncoderStep(ModularPipelineBlocks): model_name = "qwenimage" @@ -511,6 +643,61 @@ class QwenImageEditTextEncoderStep(ModularPipelineBlocks): return components, state +class QwenImageEditPlusTextEncoderStep(QwenImageEditTextEncoderStep): + model_name = "qwenimage" + + @property + def expected_configs(self) -> List[ConfigSpec]: + return [ + ConfigSpec( + name="prompt_template_encode", + default="<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n", + ), + ConfigSpec( + name="img_template_encode", + default="Picture {}: <|vision_start|><|image_pad|><|vision_end|>", + ), + ConfigSpec(name="prompt_template_encode_start_idx", default=64), + ] + + @torch.no_grad() + def __call__(self, components: QwenImageModularPipeline, state: PipelineState): + block_state = self.get_block_state(state) + + self.check_inputs(block_state.prompt, block_state.negative_prompt) + + device = components._execution_device + + block_state.prompt_embeds, block_state.prompt_embeds_mask = get_qwen_prompt_embeds_edit_plus( + components.text_encoder, + components.processor, + prompt=block_state.prompt, + image=block_state.resized_image, + prompt_template_encode=components.config.prompt_template_encode, + img_template_encode=components.config.img_template_encode, + prompt_template_encode_start_idx=components.config.prompt_template_encode_start_idx, + device=device, + ) + + if components.requires_unconditional_embeds: + negative_prompt = block_state.negative_prompt or " " + block_state.negative_prompt_embeds, block_state.negative_prompt_embeds_mask = ( + get_qwen_prompt_embeds_edit_plus( + components.text_encoder, + components.processor, + prompt=negative_prompt, + image=block_state.resized_image, + prompt_template_encode=components.config.prompt_template_encode, + img_template_encode=components.config.img_template_encode, + prompt_template_encode_start_idx=components.config.prompt_template_encode_start_idx, + device=device, + ) + ) + + self.set_block_state(state, block_state) + return components, state + + class QwenImageInpaintProcessImagesInputStep(ModularPipelineBlocks): model_name = "qwenimage" @@ -612,12 +799,7 @@ class QwenImageProcessImagesInputStep(ModularPipelineBlocks): @property def inputs(self) -> List[InputParam]: - return [ - InputParam("resized_image"), - InputParam("image"), - InputParam("height"), - InputParam("width"), - ] + return [InputParam("resized_image"), InputParam("image"), InputParam("height"), InputParam("width")] @property def intermediate_outputs(self) -> List[OutputParam]: @@ -661,6 +843,47 @@ class QwenImageProcessImagesInputStep(ModularPipelineBlocks): return components, state +class QwenImageEditPlusProcessImagesInputStep(QwenImageProcessImagesInputStep): + model_name = "qwenimage-edit-plus" + vae_image_size = 1024 * 1024 + + @property + def description(self) -> str: + return "Image Preprocess step for QwenImage Edit Plus. Unlike QwenImage Edit, QwenImage Edit Plus doesn't use the same resized image for further preprocessing." + + @property + def inputs(self) -> List[InputParam]: + return [InputParam("vae_image"), InputParam("image"), InputParam("height"), InputParam("width")] + + @torch.no_grad() + def __call__(self, components: QwenImageModularPipeline, state: PipelineState): + block_state = self.get_block_state(state) + + if block_state.vae_image is None and block_state.image is None: + raise ValueError("`vae_image` and `image` cannot be None at the same time") + + if block_state.vae_image is None: + image = block_state.image + self.check_inputs( + height=block_state.height, width=block_state.width, vae_scale_factor=components.vae_scale_factor + ) + height = block_state.height or components.default_height + width = block_state.width or components.default_width + block_state.processed_image = components.image_processor.preprocess( + image=image, height=height, width=width + ) + else: + width, height = block_state.vae_image[0].size + image = block_state.vae_image + + block_state.processed_image = components.image_processor.preprocess( + image=image, height=height, width=width + ) + + self.set_block_state(state, block_state) + return components, state + + class QwenImageVaeEncoderDynamicStep(ModularPipelineBlocks): model_name = "qwenimage" @@ -738,7 +961,6 @@ class QwenImageVaeEncoderDynamicStep(ModularPipelineBlocks): dtype=dtype, latent_channels=components.num_channels_latents, ) - setattr(block_state, self._image_latents_output_name, image_latents) self.set_block_state(state, block_state) diff --git a/src/diffusers/modular_pipelines/qwenimage/modular_blocks.py b/src/diffusers/modular_pipelines/qwenimage/modular_blocks.py index 9126766cc2..83bfcb3da4 100644 --- a/src/diffusers/modular_pipelines/qwenimage/modular_blocks.py +++ b/src/diffusers/modular_pipelines/qwenimage/modular_blocks.py @@ -37,6 +37,9 @@ from .denoise import ( ) from .encoders import ( QwenImageControlNetVaeEncoderStep, + QwenImageEditPlusProcessImagesInputStep, + QwenImageEditPlusResizeDynamicStep, + QwenImageEditPlusTextEncoderStep, QwenImageEditResizeDynamicStep, QwenImageEditTextEncoderStep, QwenImageInpaintProcessImagesInputStep, @@ -872,7 +875,151 @@ class QwenImageEditAutoBlocks(SequentialPipelineBlocks): ) -# 3. all block presets supported in QwenImage & QwenImage-Edit +#################### QwenImage Edit Plus ##################### + +# 3. QwenImage-Edit Plus + +## 3.1 QwenImage-Edit Plus / edit + +#### QwenImage-Edit Plus vl encoder: take both image and text prompts +QwenImageEditPlusVLEncoderBlocks = InsertableDict( + [ + ("resize", QwenImageEditPlusResizeDynamicStep()), + ("encode", QwenImageEditPlusTextEncoderStep()), + ] +) + + +class QwenImageEditPlusVLEncoderStep(SequentialPipelineBlocks): + model_name = "qwenimage" + block_classes = QwenImageEditPlusVLEncoderBlocks.values() + block_names = QwenImageEditPlusVLEncoderBlocks.keys() + + @property + def description(self) -> str: + return "QwenImage-Edit Plus VL encoder step that encode the image an text prompts together." + + +#### QwenImage-Edit Plus vae encoder +QwenImageEditPlusVaeEncoderBlocks = InsertableDict( + [ + ("resize", QwenImageEditPlusResizeDynamicStep()), # edit plus has a different resize step + ("preprocess", QwenImageEditPlusProcessImagesInputStep()), # vae_image -> processed_image + ("encode", QwenImageVaeEncoderDynamicStep()), # processed_image -> image_latents + ] +) + + +class QwenImageEditPlusVaeEncoderStep(SequentialPipelineBlocks): + model_name = "qwenimage" + block_classes = QwenImageEditPlusVaeEncoderBlocks.values() + block_names = QwenImageEditPlusVaeEncoderBlocks.keys() + + @property + def description(self) -> str: + return "Vae encoder step that encode the image inputs into their latent representations." + + +#### QwenImage Edit Plus presets +EDIT_PLUS_BLOCKS = InsertableDict( + [ + ("text_encoder", QwenImageEditPlusVLEncoderStep()), + ("vae_encoder", QwenImageEditPlusVaeEncoderStep()), + ("input", QwenImageEditInputStep()), + ("prepare_latents", QwenImagePrepareLatentsStep()), + ("set_timesteps", QwenImageSetTimestepsStep()), + ("prepare_rope_inputs", QwenImageEditRoPEInputsStep()), + ("denoise", QwenImageEditDenoiseStep()), + ("decode", QwenImageDecodeStep()), + ] +) + + +# auto before_denoise step for edit tasks +class QwenImageEditPlusAutoBeforeDenoiseStep(AutoPipelineBlocks): + model_name = "qwenimage-edit-plus" + block_classes = [QwenImageEditBeforeDenoiseStep] + block_names = ["edit"] + block_trigger_inputs = ["image_latents"] + + @property + def description(self): + return ( + "Before denoise step that prepare the inputs (timesteps, latents, rope inputs etc.) for the denoise step.\n" + + "This is an auto pipeline block that works for edit (img2img) task.\n" + + " - `QwenImageEditBeforeDenoiseStep` (edit) is used when `image_latents` is provided and `processed_mask_image` is not provided.\n" + + " - if `image_latents` is not provided, step will be skipped." + ) + + +## 3.2 QwenImage-Edit Plus/auto encoders + + +class QwenImageEditPlusAutoVaeEncoderStep(AutoPipelineBlocks): + block_classes = [ + QwenImageEditPlusVaeEncoderStep, + ] + block_names = ["edit"] + block_trigger_inputs = ["image"] + + @property + def description(self): + return ( + "Vae encoder step that encode the image inputs into their latent representations. \n" + " This is an auto pipeline block that works for edit task.\n" + + " - `QwenImageEditPlusVaeEncoderStep` (edit) is used when `image` is provided.\n" + + " - if `image` is not provided, step will be skipped." + ) + + +## 3.3 QwenImage-Edit/auto blocks & presets + + +class QwenImageEditPlusCoreDenoiseStep(SequentialPipelineBlocks): + model_name = "qwenimage-edit-plus" + block_classes = [ + QwenImageEditAutoInputStep, + QwenImageEditPlusAutoBeforeDenoiseStep, + QwenImageEditAutoDenoiseStep, + ] + block_names = ["input", "before_denoise", "denoise"] + + @property + def description(self): + return ( + "Core step that performs the denoising process. \n" + + " - `QwenImageEditAutoInputStep` (input) standardizes the inputs for the denoising step.\n" + + " - `QwenImageEditPlusAutoBeforeDenoiseStep` (before_denoise) prepares the inputs for the denoising step.\n" + + " - `QwenImageEditAutoDenoiseStep` (denoise) iteratively denoises the latents.\n\n" + + "This step support edit (img2img) workflow for QwenImage Edit Plus:\n" + + " - When `image_latents` is provided, it will be used for edit (img2img) task.\n" + ) + + +EDIT_PLUS_AUTO_BLOCKS = InsertableDict( + [ + ("text_encoder", QwenImageEditPlusVLEncoderStep()), + ("vae_encoder", QwenImageEditPlusAutoVaeEncoderStep()), + ("denoise", QwenImageEditPlusCoreDenoiseStep()), + ("decode", QwenImageAutoDecodeStep()), + ] +) + + +class QwenImageEditPlusAutoBlocks(SequentialPipelineBlocks): + model_name = "qwenimage-edit-plus" + block_classes = EDIT_PLUS_AUTO_BLOCKS.values() + block_names = EDIT_PLUS_AUTO_BLOCKS.keys() + + @property + def description(self): + return ( + "Auto Modular pipeline for edit (img2img) and edit tasks using QwenImage-Edit Plus.\n" + + "- for edit (img2img) generation, you need to provide `image`\n" + ) + + +# 3. all block presets supported in QwenImage, QwenImage-Edit, QwenImage-Edit Plus ALL_BLOCKS = { @@ -880,8 +1027,10 @@ ALL_BLOCKS = { "img2img": IMAGE2IMAGE_BLOCKS, "edit": EDIT_BLOCKS, "edit_inpaint": EDIT_INPAINT_BLOCKS, + "edit_plus": EDIT_PLUS_BLOCKS, "inpaint": INPAINT_BLOCKS, "controlnet": CONTROLNET_BLOCKS, "auto": AUTO_BLOCKS, "edit_auto": EDIT_AUTO_BLOCKS, + "edit_plus_auto": EDIT_PLUS_AUTO_BLOCKS, } diff --git a/src/diffusers/modular_pipelines/qwenimage/modular_pipeline.py b/src/diffusers/modular_pipelines/qwenimage/modular_pipeline.py index 7200169923..d9e30864f6 100644 --- a/src/diffusers/modular_pipelines/qwenimage/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/qwenimage/modular_pipeline.py @@ -196,3 +196,13 @@ class QwenImageEditModularPipeline(ModularPipeline, QwenImageLoraLoaderMixin): requires_unconditional_embeds = self.guider._enabled and self.guider.num_conditions > 1 return requires_unconditional_embeds + + +class QwenImageEditPlusModularPipeline(QwenImageEditModularPipeline): + """ + A ModularPipeline for QwenImage-Edit Plus. + + > [!WARNING] > This is an experimental feature and is likely to change in the future. + """ + + default_blocks_name = "QwenImageEditPlusAutoBlocks" diff --git a/src/diffusers/pipelines/auto_pipeline.py b/src/diffusers/pipelines/auto_pipeline.py index d265bfdcaf..8a32d4c367 100644 --- a/src/diffusers/pipelines/auto_pipeline.py +++ b/src/diffusers/pipelines/auto_pipeline.py @@ -95,6 +95,7 @@ from .qwenimage import ( QwenImageControlNetPipeline, QwenImageEditInpaintPipeline, QwenImageEditPipeline, + QwenImageEditPlusPipeline, QwenImageImg2ImgPipeline, QwenImageInpaintPipeline, QwenImagePipeline, @@ -186,6 +187,7 @@ AUTO_IMAGE2IMAGE_PIPELINES_MAPPING = OrderedDict( ("flux-kontext", FluxKontextPipeline), ("qwenimage", QwenImageImg2ImgPipeline), ("qwenimage-edit", QwenImageEditPipeline), + ("qwenimage-edit-plus", QwenImageEditPlusPipeline), ] ) diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index bb8fea8c8a..cf80377964 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -77,6 +77,36 @@ class QwenImageEditModularPipeline(metaclass=DummyObject): requires_backends(cls, ["torch", "transformers"]) +class QwenImageEditPlusAutoBlocks(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class QwenImageEditPlusModularPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class QwenImageModularPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] From ce90f9b2db9f459f463a3239dc4c24a04072fd43 Mon Sep 17 00:00:00 2001 From: SahilCarterr <110806554+SahilCarterr@users.noreply.github.com> Date: Mon, 6 Oct 2025 08:24:54 +0530 Subject: [PATCH 049/123] [FIX] Text to image training peft version (#12434) Fix peft error --- examples/text_to_image/requirements.txt | 2 +- examples/text_to_image/requirements_sdxl.txt | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/text_to_image/requirements.txt b/examples/text_to_image/requirements.txt index c3ffa42f0e..be05fe3fcd 100644 --- a/examples/text_to_image/requirements.txt +++ b/examples/text_to_image/requirements.txt @@ -5,4 +5,4 @@ datasets>=2.19.1 ftfy tensorboard Jinja2 -peft==0.7.0 +peft>=0.17.0 diff --git a/examples/text_to_image/requirements_sdxl.txt b/examples/text_to_image/requirements_sdxl.txt index 64cbc9205f..4dacc26ce4 100644 --- a/examples/text_to_image/requirements_sdxl.txt +++ b/examples/text_to_image/requirements_sdxl.txt @@ -5,4 +5,4 @@ ftfy tensorboard Jinja2 datasets -peft==0.7.0 \ No newline at end of file +peft>=0.17.0 \ No newline at end of file From 7f3e9b8695e80aa1f7a5ee55d025eeb8ee795602 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Mon, 6 Oct 2025 13:15:54 +0530 Subject: [PATCH 050/123] make flux ready for mellon (#12419) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * make flux ready for mellon * up * Apply suggestions from code review Co-authored-by: Álvaro Somoza --------- Co-authored-by: Álvaro Somoza --- .../modular_pipelines/flux/before_denoise.py | 4 +++ .../modular_pipelines/flux/encoders.py | 5 +++ .../modular_pipelines/flux/modular_blocks.py | 31 ++++++++++++++----- 3 files changed, 32 insertions(+), 8 deletions(-) diff --git a/src/diffusers/modular_pipelines/flux/before_denoise.py b/src/diffusers/modular_pipelines/flux/before_denoise.py index 4272066309..95858fbf6e 100644 --- a/src/diffusers/modular_pipelines/flux/before_denoise.py +++ b/src/diffusers/modular_pipelines/flux/before_denoise.py @@ -252,11 +252,13 @@ class FluxInputStep(ModularPipelineBlocks): InputParam( "prompt_embeds", required=True, + kwargs_type="denoiser_input_fields", type_hint=torch.Tensor, description="Pre-generated text embeddings. Can be generated from text_encoder step.", ), InputParam( "pooled_prompt_embeds", + kwargs_type="denoiser_input_fields", type_hint=torch.Tensor, description="Pre-generated pooled text embeddings. Can be generated from text_encoder step.", ), @@ -279,11 +281,13 @@ class FluxInputStep(ModularPipelineBlocks): OutputParam( "prompt_embeds", type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", description="text embeddings used to guide the image generation", ), OutputParam( "pooled_prompt_embeds", type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", description="pooled text embeddings used to guide the image generation", ), # TODO: support negative embeddings? diff --git a/src/diffusers/modular_pipelines/flux/encoders.py b/src/diffusers/modular_pipelines/flux/encoders.py index 8c49990280..16ddecbadb 100644 --- a/src/diffusers/modular_pipelines/flux/encoders.py +++ b/src/diffusers/modular_pipelines/flux/encoders.py @@ -181,6 +181,7 @@ class FluxTextEncoderStep(ModularPipelineBlocks): return [ InputParam("prompt"), InputParam("prompt_2"), + InputParam("max_sequence_length", type_hint=int, default=512, required=False), InputParam("joint_attention_kwargs"), ] @@ -189,16 +190,19 @@ class FluxTextEncoderStep(ModularPipelineBlocks): return [ OutputParam( "prompt_embeds", + kwargs_type="denoiser_input_fields", type_hint=torch.Tensor, description="text embeddings used to guide the image generation", ), OutputParam( "pooled_prompt_embeds", + kwargs_type="denoiser_input_fields", type_hint=torch.Tensor, description="pooled text embeddings used to guide the image generation", ), OutputParam( "text_ids", + kwargs_type="denoiser_input_fields", type_hint=torch.Tensor, description="ids from the text sequence for RoPE", ), @@ -404,6 +408,7 @@ class FluxTextEncoderStep(ModularPipelineBlocks): pooled_prompt_embeds=None, device=block_state.device, num_images_per_prompt=1, # TODO: hardcoded for now. + max_sequence_length=block_state.max_sequence_length, lora_scale=block_state.text_encoder_lora_scale, ) diff --git a/src/diffusers/modular_pipelines/flux/modular_blocks.py b/src/diffusers/modular_pipelines/flux/modular_blocks.py index 37895bddbf..ca4f993a11 100644 --- a/src/diffusers/modular_pipelines/flux/modular_blocks.py +++ b/src/diffusers/modular_pipelines/flux/modular_blocks.py @@ -84,9 +84,9 @@ class FluxImg2ImgBeforeDenoiseStep(SequentialPipelineBlocks): # before_denoise: all task (text2img, img2img) class FluxAutoBeforeDenoiseStep(AutoPipelineBlocks): - block_classes = [FluxBeforeDenoiseStep, FluxImg2ImgBeforeDenoiseStep] - block_names = ["text2image", "img2img"] - block_trigger_inputs = [None, "image_latents"] + block_classes = [FluxImg2ImgBeforeDenoiseStep, FluxBeforeDenoiseStep] + block_names = ["img2img", "text2image"] + block_trigger_inputs = ["image_latents", None] @property def description(self): @@ -124,16 +124,32 @@ class FluxAutoDecodeStep(AutoPipelineBlocks): return "Decode step that decode the denoised latents into image outputs.\n - `FluxDecodeStep`" +class FluxCoreDenoiseStep(SequentialPipelineBlocks): + block_classes = [FluxInputStep, FluxAutoBeforeDenoiseStep, FluxAutoDenoiseStep] + block_names = ["input", "before_denoise", "denoise"] + + @property + def description(self): + return ( + "Core step that performs the denoising process. \n" + + " - `FluxInputStep` (input) standardizes the inputs for the denoising step.\n" + + " - `FluxAutoBeforeDenoiseStep` (before_denoise) prepares the inputs for the denoising step.\n" + + " - `FluxAutoDenoiseStep` (denoise) iteratively denoises the latents.\n" + + "This step support text-to-image and image-to-image tasks for Flux:\n" + + " - for image-to-image generation, you need to provide `image_latents`\n" + + " - for text-to-image generation, all you need to provide is prompt embeddings" + ) + + # text2image class FluxAutoBlocks(SequentialPipelineBlocks): block_classes = [ FluxTextEncoderStep, FluxAutoVaeEncoderStep, - FluxAutoBeforeDenoiseStep, - FluxAutoDenoiseStep, + FluxCoreDenoiseStep, FluxAutoDecodeStep, ] - block_names = ["text_encoder", "image_encoder", "before_denoise", "denoise", "decoder"] + block_names = ["text_encoder", "image_encoder", "denoise", "decode"] @property def description(self): @@ -171,8 +187,7 @@ AUTO_BLOCKS = InsertableDict( [ ("text_encoder", FluxTextEncoderStep), ("image_encoder", FluxAutoVaeEncoderStep), - ("before_denoise", FluxAutoBeforeDenoiseStep), - ("denoise", FluxAutoDenoiseStep), + ("denoise", FluxCoreDenoiseStep), ("decode", FluxAutoDecodeStep), ] ) From cf4b97b2337ea0144f35ed3ec9146e27ba7a44e8 Mon Sep 17 00:00:00 2001 From: Charles Date: Mon, 6 Oct 2025 17:45:34 +0200 Subject: [PATCH 051/123] [perf] Cache version checks (#12399) --- src/diffusers/utils/import_utils.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/src/diffusers/utils/import_utils.py b/src/diffusers/utils/import_utils.py index 9399ccd2a7..97065267b0 100644 --- a/src/diffusers/utils/import_utils.py +++ b/src/diffusers/utils/import_utils.py @@ -21,6 +21,7 @@ import operator as op import os import sys from collections import OrderedDict, defaultdict +from functools import lru_cache as cache from itertools import chain from types import ModuleType from typing import Any, Tuple, Union @@ -673,6 +674,7 @@ def compare_versions(library_or_version: Union[str, Version], operation: str, re # This function was copied from: https://github.com/huggingface/accelerate/blob/874c4967d94badd24f893064cc3bef45f57cadf7/src/accelerate/utils/versions.py#L338 +@cache def is_torch_version(operation: str, version: str): """ Compares the current PyTorch version to a given reference with an operation. @@ -686,6 +688,7 @@ def is_torch_version(operation: str, version: str): return compare_versions(parse(_torch_version), operation, version) +@cache def is_torch_xla_version(operation: str, version: str): """ Compares the current torch_xla version to a given reference with an operation. @@ -701,6 +704,7 @@ def is_torch_xla_version(operation: str, version: str): return compare_versions(parse(_torch_xla_version), operation, version) +@cache def is_transformers_version(operation: str, version: str): """ Compares the current Transformers version to a given reference with an operation. @@ -716,6 +720,7 @@ def is_transformers_version(operation: str, version: str): return compare_versions(parse(_transformers_version), operation, version) +@cache def is_hf_hub_version(operation: str, version: str): """ Compares the current Hugging Face Hub version to a given reference with an operation. @@ -731,6 +736,7 @@ def is_hf_hub_version(operation: str, version: str): return compare_versions(parse(_hf_hub_version), operation, version) +@cache def is_accelerate_version(operation: str, version: str): """ Compares the current Accelerate version to a given reference with an operation. @@ -746,6 +752,7 @@ def is_accelerate_version(operation: str, version: str): return compare_versions(parse(_accelerate_version), operation, version) +@cache def is_peft_version(operation: str, version: str): """ Compares the current PEFT version to a given reference with an operation. @@ -761,6 +768,7 @@ def is_peft_version(operation: str, version: str): return compare_versions(parse(_peft_version), operation, version) +@cache def is_bitsandbytes_version(operation: str, version: str): """ Args: @@ -775,6 +783,7 @@ def is_bitsandbytes_version(operation: str, version: str): return compare_versions(parse(_bitsandbytes_version), operation, version) +@cache def is_gguf_version(operation: str, version: str): """ Compares the current Accelerate version to a given reference with an operation. @@ -790,6 +799,7 @@ def is_gguf_version(operation: str, version: str): return compare_versions(parse(_gguf_version), operation, version) +@cache def is_torchao_version(operation: str, version: str): """ Compares the current torchao version to a given reference with an operation. @@ -805,6 +815,7 @@ def is_torchao_version(operation: str, version: str): return compare_versions(parse(_torchao_version), operation, version) +@cache def is_k_diffusion_version(operation: str, version: str): """ Compares the current k-diffusion version to a given reference with an operation. @@ -820,6 +831,7 @@ def is_k_diffusion_version(operation: str, version: str): return compare_versions(parse(_k_diffusion_version), operation, version) +@cache def is_optimum_quanto_version(operation: str, version: str): """ Compares the current Accelerate version to a given reference with an operation. @@ -835,6 +847,7 @@ def is_optimum_quanto_version(operation: str, version: str): return compare_versions(parse(_optimum_quanto_version), operation, version) +@cache def is_nvidia_modelopt_version(operation: str, version: str): """ Compares the current Nvidia ModelOpt version to a given reference with an operation. @@ -850,6 +863,7 @@ def is_nvidia_modelopt_version(operation: str, version: str): return compare_versions(parse(_nvidia_modelopt_version), operation, version) +@cache def is_xformers_version(operation: str, version: str): """ Compares the current xformers version to a given reference with an operation. @@ -865,6 +879,7 @@ def is_xformers_version(operation: str, version: str): return compare_versions(parse(_xformers_version), operation, version) +@cache def is_sageattention_version(operation: str, version: str): """ Compares the current sageattention version to a given reference with an operation. @@ -880,6 +895,7 @@ def is_sageattention_version(operation: str, version: str): return compare_versions(parse(_sageattention_version), operation, version) +@cache def is_flash_attn_version(operation: str, version: str): """ Compares the current flash-attention version to a given reference with an operation. From 0974b4c6067165434fa715654b355b41beb5fceb Mon Sep 17 00:00:00 2001 From: Changseop Yeom <89627517+braintrue@users.noreply.github.com> Date: Tue, 7 Oct 2025 06:24:05 +0900 Subject: [PATCH 052/123] [i18n-KO] Fix typo and update translation in ethical_guidelines.md (#12435) --- .../ko/conceptual/ethical_guidelines.md | 42 +++++++++---------- 1 file changed, 19 insertions(+), 23 deletions(-) diff --git a/docs/source/ko/conceptual/ethical_guidelines.md b/docs/source/ko/conceptual/ethical_guidelines.md index b8c55048bf..63fc4a7741 100644 --- a/docs/source/ko/conceptual/ethical_guidelines.md +++ b/docs/source/ko/conceptual/ethical_guidelines.md @@ -14,51 +14,47 @@ specific language governing permissions and limitations under the License. ## 서문 [[preamble]] -[Diffusers](https://huggingface.co/docs/diffusers/index)는 사전 훈련된 diffusion 모델을 제공하며 추론 및 훈련을 위한 모듈식 툴박스로 사용됩니다. +[Diffusers](https://huggingface.co/docs/diffusers/index)는 사전 훈련된 diffusion 모델을 제공하며, 추론과 훈련을 위한 모듈형 툴박스로 활용됩니다. -이 기술의 실제 적용과 사회에 미칠 수 있는 부정적인 영향을 고려하여 Diffusers 라이브러리의 개발, 사용자 기여 및 사용에 윤리 지침을 제공하는 것이 중요하다고 생각합니다. - -이이 기술을 사용함에 따른 위험은 여전히 검토 중이지만, 몇 가지 예를 들면: 예술가들에 대한 저작권 문제; 딥 페이크의 악용; 부적절한 맥락에서의 성적 콘텐츠 생성; 동의 없는 사칭; 소수자 집단의 억압을 영속화하는 유해한 사회적 편견 등이 있습니다. - -우리는 위험을 지속적으로 추적하고 커뮤니티의 응답과 소중한 피드백에 따라 다음 지침을 조정할 것입니다. +이 기술의 실제 적용 사례와 사회에 미칠 수 있는 잠재적 부정적 영향을 고려할 때, Diffusers 라이브러리의 개발, 사용자 기여, 사용에 윤리 지침을 제공하는 것이 중요하다고 생각합니다. +이 기술 사용과 관련된 위험은 여전히 검토 중이지만, 예를 들면: 예술가의 저작권 문제, 딥페이크 악용, 부적절한 맥락에서의 성적 콘텐츠 생성, 비동의 사칭, 소수자 집단 억압을 영속화하는 유해한 사회적 편견 등이 있습니다. +우리는 이러한 위험을 지속적으로 추적하고, 커뮤니티의 반응과 소중한 피드백에 따라 아래 지침을 조정할 것입니다. ## 범위 [[scope]] -Diffusers 커뮤니티는 프로젝트의 개발에 다음과 같은 윤리 지침을 적용하며, 특히 윤리적 문제와 관련된 민감한 주제에 대한 커뮤니티의 기여를 조정하는 데 도움을 줄 것입니다. - +Diffusers 커뮤니티는 프로젝트 개발에 다음 윤리 지침을 적용하며, 특히 윤리적 문제와 관련된 민감한 주제에 대해 커뮤니티의 기여를 조정하는 데 도움을 줄 것입니다. ## 윤리 지침 [[ethical-guidelines]] -다음 윤리 지침은 일반적으로 적용되지만, 민감한 윤리적 문제와 관련하여 기술적 선택을 할 때 이를 우선적으로 적용할 것입니다. 나아가, 해당 기술의 최신 동향과 관련된 새로운 위험이 발생함에 따라 이러한 윤리 원칙을 조정할 것을 약속드립니다. +다음 윤리 지침은 일반적으로 적용되지만, 윤리적으로 민감한 문제와 관련된 기술적 선택을 할 때 우선적으로 적용됩니다. 또한, 해당 기술의 최신 동향과 관련된 새로운 위험이 발생함에 따라 이러한 윤리 원칙을 지속적으로 조정할 것을 약속합니다. -- **투명성**: 우리는 PR을 관리하고, 사용자에게 우리의 선택을 설명하며, 기술적 의사결정을 내릴 때 투명성을 유지할 것을 약속합니다. +- **투명성**: 우리는 PR 관리, 사용자에게 선택의 이유 설명, 기술적 의사결정 과정에서 투명성을 유지할 것을 약속합니다. -- **일관성**: 우리는 프로젝트 관리에서 사용자들에게 동일한 수준의 관심을 보장하고 기술적으로 안정되고 일관된 상태를 유지할 것을 약속합니다. +- **일관성**: 프로젝트 관리에서 모든 사용자에게 동일한 수준의 관심을 보장하고, 기술적으로 안정적이고 일관된 상태를 유지할 것을 약속합니다. -- **간결성**: Diffusers 라이브러리를 사용하고 활용하기 쉽게 만들기 위해, 프로젝트의 목표를 간결하고 일관성 있게 유지할 것을 약속합니다. +- **간결성**: Diffusers 라이브러리를 쉽게 사용하고 활용할 수 있도록, 프로젝트의 목표를 간결하고 일관성 있게 유지할 것을 약속합니다. -- **접근성**: Diffusers 프로젝트는 기술적 전문 지식 없어도 프로젝트 운영에 참여할 수 있는 기여자의 진입장벽을 낮춥니다. 이를 통해 연구 결과물이 커뮤니티에 더 잘 접근할 수 있게 됩니다. +- **접근성**: Diffusers 프로젝트는 기술적 전문지식이 없어도 기여할 수 있도록 진입장벽을 낮춥니다. 이를 통해 연구 결과물이 커뮤니티에 더 잘 접근될 수 있습니다. -- **재현성**: 우리는 Diffusers 라이브러리를 통해 제공되는 업스트림(upstream) 코드, 모델 및 데이터셋의 재현성에 대해 투명하게 공개할 것을 목표로 합니다. - -- **책임**: 우리는 커뮤니티와 팀워크를 통해, 이 기술의 잠재적인 위험과 위험을 예측하고 완화하는 데 대한 공동 책임을 가지고 있습니다. +- **재현성**: 우리는 Diffusers 라이브러리를 통해 제공되는 업스트림 코드, 모델, 데이터셋의 재현성에 대해 투명하게 공개하는 것을 목표로 합니다. +- **책임**: 커뮤니티와 팀워크를 통해, 이 기술의 잠재적 위험을 예측하고 완화하는 데 공동 책임을 집니다. ## 구현 사례: 안전 기능과 메커니즘 [[examples-of-implementations-safety-features-and-mechanisms]] -팀은 diffusion 기술과 관련된 잠재적인 윤리 및 사회적 위험에 대처하기 위한 기술적 및 비기술적 도구를 제공하고자 하고 있습니다. 또한, 커뮤니티의 참여는 이러한 기능의 구현하고 우리와 함께 인식을 높이는 데 매우 중요합니다. +팀은 diffusion 기술과 관련된 잠재적 윤리 및 사회적 위험에 대응하기 위해 기술적·비기술적 도구를 제공하고자 노력하고 있습니다. 또한, 커뮤니티의 참여는 이러한 기능 구현과 인식 제고에 매우 중요합니다. -- [**커뮤니티 탭**](https://huggingface.co/docs/hub/repositories-pull-requests-discussions): 이를 통해 커뮤니티는 프로젝트에 대해 토론하고 더 나은 협력을 할 수 있습니다. +- [**커뮤니티 탭**](https://huggingface.co/docs/hub/repositories-pull-requests-discussions): 커뮤니티가 프로젝트에 대해 토론하고 더 나은 협업을 할 수 있도록 지원합니다. -- **편향 탐색 및 평가**: Hugging Face 팀은 Stable Diffusion 모델의 편향성을 대화형으로 보여주는 [space](https://huggingface.co/spaces/society-ethics/DiffusionBiasExplorer)을 제공합니다. 이런 의미에서, 우리는 편향 탐색 및 평가를 지원하고 장려합니다. +- **편향 탐색 및 평가**: Hugging Face 팀은 Stable Diffusion 모델의 편향성을 대화형으로 보여주는 [space](https://huggingface.co/spaces/society-ethics/DiffusionBiasExplorer)를 제공합니다. 우리는 이러한 편향 탐색과 평가를 지원하고 장려합니다. - **배포에서의 안전 유도** - - [**안전한 Stable Diffusion**](https://huggingface.co/docs/diffusers/main/en/api/pipelines/stable_diffusion/stable_diffusion_safe): 이는 필터되지 않은 웹 크롤링 데이터셋으로 훈련된 Stable Diffusion과 같은 모델이 부적절한 변질에 취약한 문제를 완화합니다. 관련 논문: [Safe Latent Diffusion: Mitigating Inappropriate Degeneration in Diffusion Models](https://huggingface.co/papers/2211.05105). + - [**안전한 Stable Diffusion**](https://huggingface.co/docs/diffusers/main/en/api/pipelines/stable_diffusion/stable_diffusion_safe): 필터링되지 않은 웹 크롤링 데이터셋으로 훈련된 Stable Diffusion과 같은 모델이 부적절하게 변질되는 문제를 완화합니다. 관련 논문: [Safe Latent Diffusion: Mitigating Inappropriate Degeneration in Diffusion Models](https://huggingface.co/papers/2211.05105). - - [**안전 검사기**](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/safety_checker.py): 이미지가 생성된 후에 이미자가 임베딩 공간에서 일련의 하드코딩된 유해 개념의 클래스일 확률을 확인하고 비교합니다. 유해 개념은 역공학을 방지하기 위해 의도적으로 숨겨져 있습니다. + - [**안전 검사기**](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/safety_checker.py): 생성된 이미지가 임베딩 공간에서 하드코딩된 유해 개념 클래스와 일치할 확률을 확인하고 비교합니다. 유해 개념은 역공학을 방지하기 위해 의도적으로 숨겨져 있습니다. -- **Hub에서의 단계적인 배포**: 특히 민감한 상황에서는 일부 리포지토리에 대한 접근을 제한해야 합니다. 이 단계적인 배포는 중간 단계로, 리포지토리 작성자가 사용에 대한 더 많은 통제력을 갖게 합니다. +- **Hub에서의 단계적 배포**: 특히 민감한 상황에서는 일부 리포지토리에 대한 접근을 제한할 수 있습니다. 단계적 배포는 리포지토리 작성자가 사용에 대해 더 많은 통제권을 갖도록 하는 중간 단계입니다. -- **라이선싱**: [OpenRAILs](https://huggingface.co/blog/open_rail)와 같은 새로운 유형의 라이선싱을 통해 자유로운 접근을 보장하면서도 더 책임 있는 사용을 위한 일련의 제한을 둘 수 있습니다. +- **라이선싱**: [OpenRAILs](https://huggingface.co/blog/open_rail)와 같은 새로운 유형의 라이선스를 통해 자유로운 접근을 보장하면서도 보다 책임 있는 사용을 위한 일련의 제한을 둘 수 있습니다. From 2d69bacb007bc3e21f2cb9c7b4dd89fb71ccdcf8 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Tue, 7 Oct 2025 13:51:20 +0530 Subject: [PATCH 053/123] handle offload_state_dict when initing transformers models (#12438) --- src/diffusers/pipelines/pipeline_loading_utils.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/diffusers/pipelines/pipeline_loading_utils.py b/src/diffusers/pipelines/pipeline_loading_utils.py index b7a3e08105..dd542145d3 100644 --- a/src/diffusers/pipelines/pipeline_loading_utils.py +++ b/src/diffusers/pipelines/pipeline_loading_utils.py @@ -838,6 +838,9 @@ def load_sub_model( else: loading_kwargs["low_cpu_mem_usage"] = False + if is_transformers_model and is_transformers_version(">=", "4.57.0"): + loading_kwargs.pop("offload_state_dict") + if ( quantization_config is not None and isinstance(quantization_config, PipelineQuantizationConfig) From 1066de8c699db994ecd6beadd7d5293ffc3ead49 Mon Sep 17 00:00:00 2001 From: Linoy Tsaban <57615435+linoytsaban@users.noreply.github.com> Date: Tue, 7 Oct 2025 18:27:15 +0300 Subject: [PATCH 054/123] [Qwen LoRA training] fix bug when offloading (#12440) * fix bug when offload and cache_latents both enabled * fix bug when offload and cache_latents both enabled * fix bug when offload and cache_latents both enabled * fix bug when offload and cache_latents both enabled * fix bug when offload and cache_latents both enabled * fix bug when offload and cache_latents both enabled * fix bug when offload and cache_latents both enabled * fix bug when offload and cache_latents both enabled * fix bug when offload and cache_latents both enabled --- examples/dreambooth/train_dreambooth_lora_qwen_image.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/dreambooth/train_dreambooth_lora_qwen_image.py b/examples/dreambooth/train_dreambooth_lora_qwen_image.py index 75eae92dfb..56de160d6f 100644 --- a/examples/dreambooth/train_dreambooth_lora_qwen_image.py +++ b/examples/dreambooth/train_dreambooth_lora_qwen_image.py @@ -1338,7 +1338,7 @@ def main(args): batch["pixel_values"] = batch["pixel_values"].to( accelerator.device, non_blocking=True, dtype=vae.dtype ) - latents_cache.append(vae.encode(batch["pixel_values"]).latent_dist) + latents_cache.append(vae.encode(batch["pixel_values"]).latent_dist) if train_dataset.custom_instance_prompts: with offload_models(text_encoding_pipeline, device=accelerator.device, offload=args.offload): prompt_embeds, prompt_embeds_mask = compute_text_embeddings( From 2dc31677e12fe175950f28fd5a0c0703594e7ce4 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Wed, 8 Oct 2025 09:22:34 +0530 Subject: [PATCH 055/123] Align Flux modular more and more with Qwen modular (#12445) * start * fix * up --- .../modular_pipelines/flux/before_denoise.py | 331 ++++++------------ .../modular_pipelines/flux/denoise.py | 12 +- .../modular_pipelines/flux/encoders.py | 232 ++++++------ .../modular_pipelines/flux/inputs.py | 236 +++++++++++++ .../modular_pipelines/flux/modular_blocks.py | 185 ++++++---- 5 files changed, 574 insertions(+), 422 deletions(-) create mode 100644 src/diffusers/modular_pipelines/flux/inputs.py diff --git a/src/diffusers/modular_pipelines/flux/before_denoise.py b/src/diffusers/modular_pipelines/flux/before_denoise.py index 95858fbf6e..5f3193af0e 100644 --- a/src/diffusers/modular_pipelines/flux/before_denoise.py +++ b/src/diffusers/modular_pipelines/flux/before_denoise.py @@ -13,12 +13,12 @@ # limitations under the License. import inspect -from typing import Any, List, Optional, Tuple, Union +from typing import List, Optional, Union import numpy as np import torch -from ...models import AutoencoderKL +from ...pipelines import FluxPipeline from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import logging from ...utils.torch_utils import randn_tensor @@ -104,48 +104,6 @@ def calculate_shift( return mu -# Adapted from the original implementation. -def prepare_latents_img2img( - vae, scheduler, image, timestep, batch_size, num_channels_latents, height, width, dtype, device, generator -): - 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." - ) - - vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1) - latent_channels = vae.config.latent_channels - - # VAE applies 8x compression on images but we must also account for packing which requires - # latent height and width to be divisible by 2. - height = 2 * (int(height) // (vae_scale_factor * 2)) - width = 2 * (int(width) // (vae_scale_factor * 2)) - shape = (batch_size, num_channels_latents, height, width) - latent_image_ids = _prepare_latent_image_ids(batch_size, height // 2, width // 2, device, dtype) - - image = image.to(device=device, dtype=dtype) - if image.shape[1] != latent_channels: - image_latents = _encode_vae_image(image=image, generator=generator) - else: - image_latents = image - if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0: - # expand init_latents for batch_size - additional_image_per_prompt = batch_size // image_latents.shape[0] - image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0) - elif batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts." - ) - else: - image_latents = torch.cat([image_latents], dim=0) - - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - latents = scheduler.scale_noise(image_latents, timestep, noise) - latents = _pack_latents(latents, batch_size, num_channels_latents, height, width) - return latents, latent_image_ids - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" @@ -160,6 +118,7 @@ def retrieve_latents( raise AttributeError("Could not access latents of provided encoder_output") +# TODO: align this with Qwen patchifier def _pack_latents(latents, batch_size, num_channels_latents, height, width): latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2) latents = latents.permute(0, 2, 4, 1, 3, 5) @@ -168,35 +127,6 @@ def _pack_latents(latents, batch_size, num_channels_latents, height, width): return latents -def _prepare_latent_image_ids(batch_size, height, width, device, dtype): - latent_image_ids = torch.zeros(height, width, 3) - latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height)[:, None] - latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width)[None, :] - - latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape - - latent_image_ids = latent_image_ids.reshape( - latent_image_id_height * latent_image_id_width, latent_image_id_channels - ) - - return latent_image_ids.to(device=device, dtype=dtype) - - -# Cannot use "# Copied from" because it introduces weird indentation errors. -def _encode_vae_image(vae, image: torch.Tensor, generator: torch.Generator): - if isinstance(generator, list): - image_latents = [ - retrieve_latents(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(vae.encode(image), generator=generator) - - image_latents = (image_latents - vae.config.shift_factor) * vae.config.scaling_factor - - return image_latents - - def _get_initial_timesteps_and_optionals( transformer, scheduler, @@ -231,96 +161,6 @@ def _get_initial_timesteps_and_optionals( return timesteps, num_inference_steps, sigmas, guidance -class FluxInputStep(ModularPipelineBlocks): - model_name = "flux" - - @property - def description(self) -> str: - return ( - "Input processing step that:\n" - " 1. Determines `batch_size` and `dtype` based on `prompt_embeds`\n" - " 2. Adjusts input tensor shapes based on `batch_size` (number of prompts) and `num_images_per_prompt`\n\n" - "All input tensors are expected to have either batch_size=1 or match the batch_size\n" - "of prompt_embeds. The tensors will be duplicated across the batch dimension to\n" - "have a final batch_size of batch_size * num_images_per_prompt." - ) - - @property - def inputs(self) -> List[InputParam]: - return [ - InputParam("num_images_per_prompt", default=1), - InputParam( - "prompt_embeds", - required=True, - kwargs_type="denoiser_input_fields", - type_hint=torch.Tensor, - description="Pre-generated text embeddings. Can be generated from text_encoder step.", - ), - InputParam( - "pooled_prompt_embeds", - kwargs_type="denoiser_input_fields", - type_hint=torch.Tensor, - description="Pre-generated pooled text embeddings. Can be generated from text_encoder step.", - ), - # TODO: support negative embeddings? - ] - - @property - def intermediate_outputs(self) -> List[str]: - return [ - OutputParam( - "batch_size", - type_hint=int, - description="Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt", - ), - OutputParam( - "dtype", - type_hint=torch.dtype, - description="Data type of model tensor inputs (determined by `prompt_embeds`)", - ), - OutputParam( - "prompt_embeds", - type_hint=torch.Tensor, - kwargs_type="denoiser_input_fields", - description="text embeddings used to guide the image generation", - ), - OutputParam( - "pooled_prompt_embeds", - type_hint=torch.Tensor, - kwargs_type="denoiser_input_fields", - description="pooled text embeddings used to guide the image generation", - ), - # TODO: support negative embeddings? - ] - - def check_inputs(self, components, block_state): - if block_state.prompt_embeds is not None and block_state.pooled_prompt_embeds is not None: - if block_state.prompt_embeds.shape[0] != block_state.pooled_prompt_embeds.shape[0]: - raise ValueError( - "`prompt_embeds` and `pooled_prompt_embeds` must have the same batch size when passed directly, but" - f" got: `prompt_embeds` {block_state.prompt_embeds.shape} != `pooled_prompt_embeds`" - f" {block_state.pooled_prompt_embeds.shape}." - ) - - @torch.no_grad() - def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState: - # TODO: consider adding negative embeddings? - block_state = self.get_block_state(state) - self.check_inputs(components, block_state) - - block_state.batch_size = block_state.prompt_embeds.shape[0] - block_state.dtype = block_state.prompt_embeds.dtype - - _, seq_len, _ = block_state.prompt_embeds.shape - block_state.prompt_embeds = block_state.prompt_embeds.repeat(1, block_state.num_images_per_prompt, 1) - block_state.prompt_embeds = block_state.prompt_embeds.view( - block_state.batch_size * block_state.num_images_per_prompt, seq_len, -1 - ) - self.set_block_state(state, block_state) - - return components, state - - class FluxSetTimestepsStep(ModularPipelineBlocks): model_name = "flux" @@ -389,6 +229,10 @@ class FluxSetTimestepsStep(ModularPipelineBlocks): block_state.sigmas = sigmas block_state.guidance = guidance + # We set the index here to remove DtoH sync, helpful especially during compilation. + # Check out more details here: https://github.com/huggingface/diffusers/pull/11696 + components.scheduler.set_begin_index(0) + self.set_block_state(state, block_state) return components, state @@ -432,11 +276,6 @@ class FluxImg2ImgSetTimestepsStep(ModularPipelineBlocks): type_hint=int, description="The number of denoising steps to perform at inference time", ), - OutputParam( - "latent_timestep", - type_hint=torch.Tensor, - description="The timestep that represents the initial noise level for image-to-image generation", - ), OutputParam("guidance", type_hint=torch.Tensor, description="Optional guidance to be used."), ] @@ -484,8 +323,6 @@ class FluxImg2ImgSetTimestepsStep(ModularPipelineBlocks): block_state.sigmas = sigmas block_state.guidance = guidance - block_state.latent_timestep = timesteps[:1].repeat(batch_size) - self.set_block_state(state, block_state) return components, state @@ -524,11 +361,6 @@ class FluxPrepareLatentsStep(ModularPipelineBlocks): OutputParam( "latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process" ), - OutputParam( - "latent_image_ids", - type_hint=torch.Tensor, - description="IDs computed from the image sequence needed for RoPE", - ), ] @staticmethod @@ -552,20 +384,13 @@ class FluxPrepareLatentsStep(ModularPipelineBlocks): generator, latents=None, ): - # Couldn't use the `prepare_latents` method directly from Flux because I decided to copy over - # the packing methods here. So, for example, `comp._pack_latents()` won't work if we were - # to go with the "# Copied from ..." approach. Or maybe there's a way? - - # VAE applies 8x compression on images but we must also account for packing which requires - # latent height and width to be divisible by 2. height = 2 * (int(height) // (comp.vae_scale_factor * 2)) width = 2 * (int(width) // (comp.vae_scale_factor * 2)) shape = (batch_size, num_channels_latents, height, width) if latents is not None: - latent_image_ids = _prepare_latent_image_ids(batch_size, height // 2, width // 2, device, dtype) - return latents.to(device=device, dtype=dtype), latent_image_ids + return latents.to(device=device, dtype=dtype) if isinstance(generator, list) and len(generator) != batch_size: raise ValueError( @@ -573,12 +398,11 @@ class FluxPrepareLatentsStep(ModularPipelineBlocks): f" size of {batch_size}. Make sure the batch size matches the length of the generators." ) + # TODO: move packing latents code to a patchifier latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) latents = _pack_latents(latents, batch_size, num_channels_latents, height, width) - latent_image_ids = _prepare_latent_image_ids(batch_size, height // 2, width // 2, device, dtype) - - return latents, latent_image_ids + return latents @torch.no_grad() def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState: @@ -587,12 +411,11 @@ class FluxPrepareLatentsStep(ModularPipelineBlocks): block_state.height = block_state.height or components.default_height block_state.width = block_state.width or components.default_width block_state.device = components._execution_device - block_state.dtype = torch.bfloat16 # TODO: okay to hardcode this? block_state.num_channels_latents = components.num_channels_latents self.check_inputs(components, block_state) batch_size = block_state.batch_size * block_state.num_images_per_prompt - block_state.latents, block_state.latent_image_ids = self.prepare_latents( + block_state.latents = self.prepare_latents( components, batch_size, block_state.num_channels_latents, @@ -612,82 +435,124 @@ class FluxPrepareLatentsStep(ModularPipelineBlocks): class FluxImg2ImgPrepareLatentsStep(ModularPipelineBlocks): model_name = "flux" - @property - def expected_components(self) -> List[ComponentSpec]: - return [ComponentSpec("vae", AutoencoderKL), ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] - @property def description(self) -> str: - return "Step that prepares the latents for the image-to-image generation process" + return "Step that adds noise to image latents for image-to-image. Should be run after `set_timesteps`," + " `prepare_latents`. Both noise and image latents should already be patchified." @property - def inputs(self) -> List[Tuple[str, Any]]: + def expected_components(self) -> List[ComponentSpec]: + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def inputs(self) -> List[InputParam]: return [ - InputParam("height", type_hint=int), - InputParam("width", type_hint=int), - InputParam("latents", type_hint=Optional[torch.Tensor]), - InputParam("num_images_per_prompt", type_hint=int, default=1), - InputParam("generator"), InputParam( - "image_latents", + name="latents", required=True, type_hint=torch.Tensor, - description="The latents representing the reference image for image-to-image/inpainting generation. Can be generated in vae_encode step.", + description="The initial random noised, can be generated in prepare latent step.", ), InputParam( - "latent_timestep", + name="image_latents", required=True, type_hint=torch.Tensor, - description="The timestep that represents the initial noise level for image-to-image/inpainting generation. Can be generated in set_timesteps step.", + description="The image latents to use for the denoising process. Can be generated in vae encoder and packed in input step.", ), InputParam( - "batch_size", + name="timesteps", required=True, - type_hint=int, - description="Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can be generated in input step.", + type_hint=torch.Tensor, + description="The timesteps to use for the denoising process. Can be generated in set_timesteps step.", ), - InputParam("dtype", required=True, type_hint=torch.dtype, description="The dtype of the model inputs"), ] @property def intermediate_outputs(self) -> List[OutputParam]: return [ OutputParam( - "latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process" - ), - OutputParam( - "latent_image_ids", + name="initial_noise", type_hint=torch.Tensor, - description="IDs computed from the image sequence needed for RoPE", + description="The initial random noised used for inpainting denoising.", ), ] + @staticmethod + def check_inputs(image_latents, latents): + if image_latents.shape[0] != latents.shape[0]: + raise ValueError( + f"`image_latents` must have have same batch size as `latents`, but got {image_latents.shape[0]} and {latents.shape[0]}" + ) + + if image_latents.ndim != 3: + raise ValueError(f"`image_latents` must have 3 dimensions (patchified), but got {image_latents.ndim}") + @torch.no_grad() def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) - block_state.device = components._execution_device - block_state.dtype = torch.bfloat16 # TODO: okay to hardcode this? - block_state.num_channels_latents = components.num_channels_latents - block_state.dtype = block_state.dtype if block_state.dtype is not None else components.vae.dtype - block_state.device = components._execution_device + self.check_inputs(image_latents=block_state.image_latents, latents=block_state.latents) - # TODO: implement `check_inputs` - batch_size = block_state.batch_size * block_state.num_images_per_prompt - if block_state.latents is None: - block_state.latents, block_state.latent_image_ids = prepare_latents_img2img( - components.vae, - components.scheduler, - block_state.image_latents, - block_state.latent_timestep, - batch_size, - block_state.num_channels_latents, - block_state.height, - block_state.width, - block_state.dtype, - block_state.device, - block_state.generator, - ) + # prepare latent timestep + latent_timestep = block_state.timesteps[:1].repeat(block_state.latents.shape[0]) + + # make copy of initial_noise + block_state.initial_noise = block_state.latents + + # scale noise + block_state.latents = components.scheduler.scale_noise( + block_state.image_latents, latent_timestep, block_state.latents + ) + + self.set_block_state(state, block_state) + + return components, state + + +class FluxRoPEInputsStep(ModularPipelineBlocks): + model_name = "flux" + + @property + def description(self) -> str: + return "Step that prepares the RoPE inputs for the denoising process. Should be placed after text encoder and latent preparation steps." + + @property + def inputs(self) -> List[InputParam]: + return [ + InputParam(name="height", required=True), + InputParam(name="width", required=True), + InputParam(name="prompt_embeds"), + ] + + @property + def intermediate_outputs(self) -> List[OutputParam]: + return [ + OutputParam( + name="txt_ids", + kwargs_type="denoiser_input_fields", + type_hint=List[int], + description="The sequence lengths of the prompt embeds, used for RoPE calculation.", + ), + OutputParam( + name="img_ids", + kwargs_type="denoiser_input_fields", + type_hint=List[int], + description="The sequence lengths of the image latents, used for RoPE calculation.", + ), + ] + + def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + prompt_embeds = block_state.prompt_embeds + device, dtype = prompt_embeds.device, prompt_embeds.dtype + block_state.txt_ids = torch.zeros(prompt_embeds.shape[1], 3).to( + device=prompt_embeds.device, dtype=prompt_embeds.dtype + ) + + height = 2 * (int(block_state.height) // (components.vae_scale_factor * 2)) + width = 2 * (int(block_state.width) // (components.vae_scale_factor * 2)) + block_state.img_ids = FluxPipeline._prepare_latent_image_ids(None, height // 2, width // 2, device, dtype) self.set_block_state(state, block_state) diff --git a/src/diffusers/modular_pipelines/flux/denoise.py b/src/diffusers/modular_pipelines/flux/denoise.py index ffa0a4456f..e482c198e8 100644 --- a/src/diffusers/modular_pipelines/flux/denoise.py +++ b/src/diffusers/modular_pipelines/flux/denoise.py @@ -76,18 +76,17 @@ class FluxLoopDenoiser(ModularPipelineBlocks): description="Pooled prompt embeddings", ), InputParam( - "text_ids", + "txt_ids", required=True, type_hint=torch.Tensor, description="IDs computed from text sequence needed for RoPE", ), InputParam( - "latent_image_ids", + "img_ids", required=True, type_hint=torch.Tensor, description="IDs computed from image sequence needed for RoPE", ), - # TODO: guidance ] @torch.no_grad() @@ -101,8 +100,8 @@ class FluxLoopDenoiser(ModularPipelineBlocks): encoder_hidden_states=block_state.prompt_embeds, pooled_projections=block_state.pooled_prompt_embeds, joint_attention_kwargs=block_state.joint_attention_kwargs, - txt_ids=block_state.text_ids, - img_ids=block_state.latent_image_ids, + txt_ids=block_state.txt_ids, + img_ids=block_state.img_ids, return_dict=False, )[0] block_state.noise_pred = noise_pred @@ -195,9 +194,6 @@ class FluxDenoiseLoopWrapper(LoopSequentialPipelineBlocks): block_state.num_warmup_steps = max( len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order, 0 ) - # We set the index here to remove DtoH sync, helpful especially during compilation. - # Check out more details here: https://github.com/huggingface/diffusers/pull/11696 - components.scheduler.set_begin_index(0) with self.progress_bar(total=block_state.num_inference_steps) as progress_bar: for i, t in enumerate(block_state.timesteps): components, block_state = self.loop_step(components, block_state, i=i, t=t) diff --git a/src/diffusers/modular_pipelines/flux/encoders.py b/src/diffusers/modular_pipelines/flux/encoders.py index 16ddecbadb..6368086cbb 100644 --- a/src/diffusers/modular_pipelines/flux/encoders.py +++ b/src/diffusers/modular_pipelines/flux/encoders.py @@ -25,7 +25,7 @@ from ...loaders import FluxLoraLoaderMixin, TextualInversionLoaderMixin from ...models import AutoencoderKL from ...utils import USE_PEFT_BACKEND, is_ftfy_available, logging, scale_lora_layers, unscale_lora_layers from ..modular_pipeline import ModularPipelineBlocks, PipelineState -from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam from .modular_pipeline import FluxModularPipeline @@ -67,89 +67,148 @@ def retrieve_latents( raise AttributeError("Could not access latents of provided encoder_output") -class FluxVaeEncoderStep(ModularPipelineBlocks): - model_name = "flux" +def encode_vae_image(vae: AutoencoderKL, image: torch.Tensor, generator: torch.Generator, sample_mode="sample"): + if isinstance(generator, list): + image_latents = [ + retrieve_latents(vae.encode(image[i : i + 1]), generator=generator[i], sample_mode=sample_mode) + for i in range(image.shape[0]) + ] + image_latents = torch.cat(image_latents, dim=0) + else: + image_latents = retrieve_latents(vae.encode(image), generator=generator, sample_mode=sample_mode) + + image_latents = (image_latents - vae.config.shift_factor) * vae.config.scaling_factor + + return image_latents + + +class FluxProcessImagesInputStep(ModularPipelineBlocks): + model_name = "Flux" @property def description(self) -> str: - return "Vae Encoder step that encode the input image into a latent representation" + return "Image Preprocess step. Resizing is needed in Flux Kontext (will be implemented later.)" @property def expected_components(self) -> List[ComponentSpec]: return [ - ComponentSpec("vae", AutoencoderKL), ComponentSpec( "image_processor", VaeImageProcessor, - config=FrozenDict({"vae_scale_factor": 16, "vae_latent_channels": 16}), + config=FrozenDict({"vae_scale_factor": 16}), default_creation_method="from_config", ), ] @property def inputs(self) -> List[InputParam]: + return [InputParam("resized_image"), InputParam("image"), InputParam("height"), InputParam("width")] + + @property + def intermediate_outputs(self) -> List[OutputParam]: return [ - InputParam("image", required=True), - InputParam("height"), - InputParam("width"), - InputParam("generator"), - InputParam("dtype", type_hint=torch.dtype, description="Data type of model tensor inputs"), - InputParam( - "preprocess_kwargs", - type_hint=Optional[dict], - description="A kwargs dictionary that if specified is passed along to the `ImageProcessor` as defined under `self.image_processor` in [diffusers.image_processor.VaeImageProcessor]", - ), + OutputParam(name="processed_image"), ] + @staticmethod + def check_inputs(height, width, vae_scale_factor): + if height is not None and height % (vae_scale_factor * 2) != 0: + raise ValueError(f"Height must be divisible by {vae_scale_factor * 2} but is {height}") + + if width is not None and width % (vae_scale_factor * 2) != 0: + raise ValueError(f"Width must be divisible by {vae_scale_factor * 2} but is {width}") + + @torch.no_grad() + def __call__(self, components: FluxModularPipeline, state: PipelineState): + block_state = self.get_block_state(state) + + if block_state.resized_image is None and block_state.image is None: + raise ValueError("`resized_image` and `image` cannot be None at the same time") + + if block_state.resized_image is None: + image = block_state.image + self.check_inputs( + height=block_state.height, width=block_state.width, vae_scale_factor=components.vae_scale_factor + ) + height = block_state.height or components.default_height + width = block_state.width or components.default_width + else: + width, height = block_state.resized_image[0].size + image = block_state.resized_image + + block_state.processed_image = components.image_processor.preprocess(image=image, height=height, width=width) + + self.set_block_state(state, block_state) + return components, state + + +class FluxVaeEncoderDynamicStep(ModularPipelineBlocks): + model_name = "flux" + + def __init__( + self, + input_name: str = "processed_image", + output_name: str = "image_latents", + ): + """Initialize a VAE encoder step for converting images to latent representations. + + Both the input and output names are configurable so this block can be configured to process to different image + inputs (e.g., "processed_image" -> "image_latents", "processed_control_image" -> "control_image_latents"). + + Args: + input_name (str, optional): Name of the input image tensor. Defaults to "processed_image". + Examples: "processed_image" or "processed_control_image" + output_name (str, optional): Name of the output latent tensor. Defaults to "image_latents". + Examples: "image_latents" or "control_image_latents" + + Examples: + # Basic usage with default settings (includes image processor): # FluxImageVaeEncoderDynamicStep() + + # Custom input/output names for control image: # FluxImageVaeEncoderDynamicStep( + input_name="processed_control_image", output_name="control_image_latents" + ) + """ + self._image_input_name = input_name + self._image_latents_output_name = output_name + super().__init__() + + @property + def description(self) -> str: + return f"Dynamic VAE Encoder step that converts {self._image_input_name} into latent representations {self._image_latents_output_name}.\n" + + @property + def expected_components(self) -> List[ComponentSpec]: + components = [ComponentSpec("vae", AutoencoderKL)] + return components + + @property + def inputs(self) -> List[InputParam]: + inputs = [InputParam(self._image_input_name, required=True), InputParam("generator")] + return inputs + @property def intermediate_outputs(self) -> List[OutputParam]: return [ OutputParam( - "image_latents", + self._image_latents_output_name, type_hint=torch.Tensor, - description="The latents representing the reference image for image-to-image/inpainting generation", + description="The latents representing the reference image", ) ] - @staticmethod - # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3_inpaint.StableDiffusion3InpaintPipeline._encode_vae_image with self.vae->vae - def _encode_vae_image(vae, image: torch.Tensor, generator: torch.Generator): - if isinstance(generator, list): - image_latents = [ - retrieve_latents(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(vae.encode(image), generator=generator) - - image_latents = (image_latents - vae.config.shift_factor) * vae.config.scaling_factor - - return image_latents - @torch.no_grad() def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) - block_state.preprocess_kwargs = block_state.preprocess_kwargs or {} - block_state.device = components._execution_device - block_state.dtype = block_state.dtype if block_state.dtype is not None else components.vae.dtype - block_state.image = components.image_processor.preprocess( - block_state.image, height=block_state.height, width=block_state.width, **block_state.preprocess_kwargs - ) - block_state.image = block_state.image.to(device=block_state.device, dtype=block_state.dtype) + device = components._execution_device + dtype = components.vae.dtype - block_state.batch_size = block_state.image.shape[0] + image = getattr(block_state, self._image_input_name) + image = image.to(device=device, dtype=dtype) - # if generator is a list, make sure the length of it matches the length of images (both should be batch_size) - if isinstance(block_state.generator, list) and len(block_state.generator) != block_state.batch_size: - raise ValueError( - f"You have passed a list of generators of length {len(block_state.generator)}, but requested an effective batch" - f" size of {block_state.batch_size}. Make sure the batch size matches the length of the generators." - ) - - block_state.image_latents = self._encode_vae_image( - components.vae, image=block_state.image, generator=block_state.generator - ) + # Encode image into latents + image_latents = encode_vae_image(image=image, vae=components.vae, generator=block_state.generator) + setattr(block_state, self._image_latents_output_name, image_latents) self.set_block_state(state, block_state) @@ -161,7 +220,7 @@ class FluxTextEncoderStep(ModularPipelineBlocks): @property def description(self) -> str: - return "Text Encoder step that generate text_embeddings to guide the video generation" + return "Text Encoder step that generate text_embeddings to guide the image generation" @property def expected_components(self) -> List[ComponentSpec]: @@ -172,10 +231,6 @@ class FluxTextEncoderStep(ModularPipelineBlocks): ComponentSpec("tokenizer_2", T5TokenizerFast), ] - @property - def expected_configs(self) -> List[ConfigSpec]: - return [] - @property def inputs(self) -> List[InputParam]: return [ @@ -200,12 +255,6 @@ class FluxTextEncoderStep(ModularPipelineBlocks): type_hint=torch.Tensor, description="pooled text embeddings used to guide the image generation", ), - OutputParam( - "text_ids", - kwargs_type="denoiser_input_fields", - type_hint=torch.Tensor, - description="ids from the text sequence for RoPE", - ), ] @staticmethod @@ -216,16 +265,10 @@ class FluxTextEncoderStep(ModularPipelineBlocks): @staticmethod def _get_t5_prompt_embeds( - components, - prompt: Union[str, List[str]], - num_images_per_prompt: int, - max_sequence_length: int, - device: torch.device, + components, prompt: Union[str, List[str]], max_sequence_length: int, device: torch.device ): dtype = components.text_encoder_2.dtype - prompt = [prompt] if isinstance(prompt, str) else prompt - batch_size = len(prompt) if isinstance(components, TextualInversionLoaderMixin): prompt = components.maybe_convert_prompt(prompt, components.tokenizer_2) @@ -251,23 +294,11 @@ class FluxTextEncoderStep(ModularPipelineBlocks): prompt_embeds = components.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0] prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) - _, seq_len, _ = prompt_embeds.shape - - # duplicate text embeddings and attention mask 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(batch_size * num_images_per_prompt, seq_len, -1) - return prompt_embeds @staticmethod - def _get_clip_prompt_embeds( - components, - prompt: Union[str, List[str]], - num_images_per_prompt: int, - device: torch.device, - ): + def _get_clip_prompt_embeds(components, prompt: Union[str, List[str]], device: torch.device): prompt = [prompt] if isinstance(prompt, str) else prompt - batch_size = len(prompt) if isinstance(components, TextualInversionLoaderMixin): prompt = components.maybe_convert_prompt(prompt, components.tokenizer) @@ -297,10 +328,6 @@ class FluxTextEncoderStep(ModularPipelineBlocks): prompt_embeds = prompt_embeds.pooler_output prompt_embeds = prompt_embeds.to(dtype=components.text_encoder.dtype, device=device) - # duplicate text embeddings for each generation per prompt, using mps friendly method - prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt) - prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1) - return prompt_embeds @staticmethod @@ -309,34 +336,11 @@ class FluxTextEncoderStep(ModularPipelineBlocks): prompt: Union[str, List[str]], prompt_2: Union[str, List[str]], device: Optional[torch.device] = None, - num_images_per_prompt: int = 1, prompt_embeds: Optional[torch.FloatTensor] = None, pooled_prompt_embeds: Optional[torch.FloatTensor] = None, max_sequence_length: int = 512, lora_scale: Optional[float] = 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 all text-encoders - device: (`torch.device`): - torch device - num_images_per_prompt (`int`): - number of images that should be generated per prompt - prompt_embeds (`torch.FloatTensor`, *optional*): - Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not - provided, text embeddings will be generated from `prompt` input argument. - pooled_prompt_embeds (`torch.FloatTensor`, *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. - lora_scale (`float`, *optional*): - A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded. - """ device = device or components._execution_device # set lora scale so that monkey patched LoRA @@ -361,12 +365,10 @@ class FluxTextEncoderStep(ModularPipelineBlocks): components, prompt=prompt, device=device, - num_images_per_prompt=num_images_per_prompt, ) prompt_embeds = FluxTextEncoderStep._get_t5_prompt_embeds( components, prompt=prompt_2, - num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, device=device, ) @@ -381,10 +383,7 @@ class FluxTextEncoderStep(ModularPipelineBlocks): # Retrieve the original scale by scaling back the LoRA layers unscale_lora_layers(components.text_encoder_2, lora_scale) - dtype = components.text_encoder.dtype if components.text_encoder is not None else torch.bfloat16 - text_ids = torch.zeros(prompt_embeds.shape[1], 3).to(device=device, dtype=dtype) - - return prompt_embeds, pooled_prompt_embeds, text_ids + return prompt_embeds, pooled_prompt_embeds @torch.no_grad() def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState: @@ -400,14 +399,13 @@ class FluxTextEncoderStep(ModularPipelineBlocks): if block_state.joint_attention_kwargs is not None else None ) - (block_state.prompt_embeds, block_state.pooled_prompt_embeds, block_state.text_ids) = self.encode_prompt( + block_state.prompt_embeds, block_state.pooled_prompt_embeds = self.encode_prompt( components, prompt=block_state.prompt, prompt_2=None, prompt_embeds=None, pooled_prompt_embeds=None, device=block_state.device, - num_images_per_prompt=1, # TODO: hardcoded for now. max_sequence_length=block_state.max_sequence_length, lora_scale=block_state.text_encoder_lora_scale, ) diff --git a/src/diffusers/modular_pipelines/flux/inputs.py b/src/diffusers/modular_pipelines/flux/inputs.py new file mode 100644 index 0000000000..f9192655d1 --- /dev/null +++ b/src/diffusers/modular_pipelines/flux/inputs.py @@ -0,0 +1,236 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import List + +import torch + +from ...pipelines import FluxPipeline +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import InputParam, OutputParam + +# TODO: consider making these common utilities for modular if they are not pipeline-specific. +from ..qwenimage.inputs import calculate_dimension_from_latents, repeat_tensor_to_batch_size +from .modular_pipeline import FluxModularPipeline + + +class FluxTextInputStep(ModularPipelineBlocks): + model_name = "flux" + + @property + def description(self) -> str: + return ( + "Text input processing step that standardizes text embeddings for the pipeline.\n" + "This step:\n" + " 1. Determines `batch_size` and `dtype` based on `prompt_embeds`\n" + " 2. Ensures all text embeddings have consistent batch sizes (batch_size * num_images_per_prompt)" + ) + + @property + def inputs(self) -> List[InputParam]: + return [ + InputParam("num_images_per_prompt", default=1), + InputParam( + "prompt_embeds", + required=True, + kwargs_type="denoiser_input_fields", + type_hint=torch.Tensor, + description="Pre-generated text embeddings. Can be generated from text_encoder step.", + ), + InputParam( + "pooled_prompt_embeds", + kwargs_type="denoiser_input_fields", + type_hint=torch.Tensor, + description="Pre-generated pooled text embeddings. Can be generated from text_encoder step.", + ), + # TODO: support negative embeddings? + ] + + @property + def intermediate_outputs(self) -> List[str]: + return [ + OutputParam( + "batch_size", + type_hint=int, + description="Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt", + ), + OutputParam( + "dtype", + type_hint=torch.dtype, + description="Data type of model tensor inputs (determined by `prompt_embeds`)", + ), + OutputParam( + "prompt_embeds", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description="text embeddings used to guide the image generation", + ), + OutputParam( + "pooled_prompt_embeds", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description="pooled text embeddings used to guide the image generation", + ), + # TODO: support negative embeddings? + ] + + def check_inputs(self, components, block_state): + if block_state.prompt_embeds is not None and block_state.pooled_prompt_embeds is not None: + if block_state.prompt_embeds.shape[0] != block_state.pooled_prompt_embeds.shape[0]: + raise ValueError( + "`prompt_embeds` and `pooled_prompt_embeds` must have the same batch size when passed directly, but" + f" got: `prompt_embeds` {block_state.prompt_embeds.shape} != `pooled_prompt_embeds`" + f" {block_state.pooled_prompt_embeds.shape}." + ) + + @torch.no_grad() + def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState: + # TODO: consider adding negative embeddings? + block_state = self.get_block_state(state) + self.check_inputs(components, block_state) + + block_state.batch_size = block_state.prompt_embeds.shape[0] + block_state.dtype = block_state.prompt_embeds.dtype + + _, seq_len, _ = block_state.prompt_embeds.shape + block_state.prompt_embeds = block_state.prompt_embeds.repeat(1, block_state.num_images_per_prompt, 1) + block_state.prompt_embeds = block_state.prompt_embeds.view( + block_state.batch_size * block_state.num_images_per_prompt, seq_len, -1 + ) + self.set_block_state(state, block_state) + + return components, state + + +# Adapted from `QwenImageInputsDynamicStep` +class FluxInputsDynamicStep(ModularPipelineBlocks): + model_name = "flux" + + def __init__( + self, + image_latent_inputs: List[str] = ["image_latents"], + additional_batch_inputs: List[str] = [], + ): + if not isinstance(image_latent_inputs, list): + image_latent_inputs = [image_latent_inputs] + if not isinstance(additional_batch_inputs, list): + additional_batch_inputs = [additional_batch_inputs] + + self._image_latent_inputs = image_latent_inputs + self._additional_batch_inputs = additional_batch_inputs + super().__init__() + + @property + def description(self) -> str: + # Functionality section + summary_section = ( + "Input processing step that:\n" + " 1. For image latent inputs: Updates height/width if None, patchifies latents, and expands batch size\n" + " 2. For additional batch inputs: Expands batch dimensions to match final batch size" + ) + + # Inputs info + inputs_info = "" + if self._image_latent_inputs or self._additional_batch_inputs: + inputs_info = "\n\nConfigured inputs:" + if self._image_latent_inputs: + inputs_info += f"\n - Image latent inputs: {self._image_latent_inputs}" + if self._additional_batch_inputs: + inputs_info += f"\n - Additional batch inputs: {self._additional_batch_inputs}" + + # Placement guidance + placement_section = "\n\nThis block should be placed after the encoder steps and the text input step." + + return summary_section + inputs_info + placement_section + + @property + def inputs(self) -> List[InputParam]: + inputs = [ + InputParam(name="num_images_per_prompt", default=1), + InputParam(name="batch_size", required=True), + InputParam(name="height"), + InputParam(name="width"), + ] + + # Add image latent inputs + for image_latent_input_name in self._image_latent_inputs: + inputs.append(InputParam(name=image_latent_input_name)) + + # Add additional batch inputs + for input_name in self._additional_batch_inputs: + inputs.append(InputParam(name=input_name)) + + return inputs + + @property + def intermediate_outputs(self) -> List[OutputParam]: + return [ + OutputParam(name="image_height", type_hint=int, description="The height of the image latents"), + OutputParam(name="image_width", type_hint=int, description="The width of the image latents"), + ] + + def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + # Process image latent inputs (height/width calculation, patchify, and batch expansion) + for image_latent_input_name in self._image_latent_inputs: + image_latent_tensor = getattr(block_state, image_latent_input_name) + if image_latent_tensor is None: + continue + + # 1. Calculate height/width from latents + height, width = calculate_dimension_from_latents(image_latent_tensor, components.vae_scale_factor) + block_state.height = block_state.height or height + block_state.width = block_state.width or width + + if not hasattr(block_state, "image_height"): + block_state.image_height = height + if not hasattr(block_state, "image_width"): + block_state.image_width = width + + # 2. Patchify the image latent tensor + # TODO: Implement patchifier for Flux. + latent_height, latent_width = image_latent_tensor.shape[2:] + image_latent_tensor = FluxPipeline._pack_latents( + image_latent_tensor, block_state.batch_size, image_latent_tensor.shape[1], latent_height, latent_width + ) + + # 3. Expand batch size + image_latent_tensor = repeat_tensor_to_batch_size( + input_name=image_latent_input_name, + input_tensor=image_latent_tensor, + num_images_per_prompt=block_state.num_images_per_prompt, + batch_size=block_state.batch_size, + ) + + setattr(block_state, image_latent_input_name, image_latent_tensor) + + # Process additional batch inputs (only batch expansion) + for input_name in self._additional_batch_inputs: + input_tensor = getattr(block_state, input_name) + if input_tensor is None: + continue + + # Only expand batch size + input_tensor = repeat_tensor_to_batch_size( + input_name=input_name, + input_tensor=input_tensor, + num_images_per_prompt=block_state.num_images_per_prompt, + batch_size=block_state.batch_size, + ) + + setattr(block_state, input_name, input_tensor) + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/flux/modular_blocks.py b/src/diffusers/modular_pipelines/flux/modular_blocks.py index ca4f993a11..b40dfe1762 100644 --- a/src/diffusers/modular_pipelines/flux/modular_blocks.py +++ b/src/diffusers/modular_pipelines/flux/modular_blocks.py @@ -18,21 +18,41 @@ from ..modular_pipeline_utils import InsertableDict from .before_denoise import ( FluxImg2ImgPrepareLatentsStep, FluxImg2ImgSetTimestepsStep, - FluxInputStep, FluxPrepareLatentsStep, + FluxRoPEInputsStep, FluxSetTimestepsStep, ) from .decoders import FluxDecodeStep from .denoise import FluxDenoiseStep -from .encoders import FluxTextEncoderStep, FluxVaeEncoderStep +from .encoders import FluxProcessImagesInputStep, FluxTextEncoderStep, FluxVaeEncoderDynamicStep +from .inputs import FluxInputsDynamicStep, FluxTextInputStep logger = logging.get_logger(__name__) # pylint: disable=invalid-name # vae encoder (run before before_denoise) +FluxImg2ImgVaeEncoderBlocks = InsertableDict( + [ + ("preprocess", FluxProcessImagesInputStep()), + ("encode", FluxVaeEncoderDynamicStep()), + ] +) + + +class FluxImg2ImgVaeEncoderStep(SequentialPipelineBlocks): + model_name = "flux" + + block_classes = FluxImg2ImgVaeEncoderBlocks.values() + block_names = FluxImg2ImgVaeEncoderBlocks.keys() + + @property + def description(self) -> str: + return "Vae encoder step that preprocess andencode the image inputs into their latent representations." + + class FluxAutoVaeEncoderStep(AutoPipelineBlocks): - block_classes = [FluxVaeEncoderStep] + block_classes = [FluxImg2ImgVaeEncoderStep] block_names = ["img2img"] block_trigger_inputs = ["image"] @@ -41,45 +61,48 @@ class FluxAutoVaeEncoderStep(AutoPipelineBlocks): return ( "Vae encoder step that encode the image inputs into their latent representations.\n" + "This is an auto pipeline block that works for img2img tasks.\n" - + " - `FluxVaeEncoderStep` (img2img) is used when only `image` is provided." - + " - if `image` is provided, step will be skipped." + + " - `FluxImg2ImgVaeEncoderStep` (img2img) is used when only `image` is provided." + + " - if `image` is not provided, step will be skipped." ) -# before_denoise: text2img, img2img -class FluxBeforeDenoiseStep(SequentialPipelineBlocks): - block_classes = [ - FluxInputStep, - FluxPrepareLatentsStep, - FluxSetTimestepsStep, +# before_denoise: text2img +FluxBeforeDenoiseBlocks = InsertableDict( + [ + ("prepare_latents", FluxPrepareLatentsStep()), + ("set_timesteps", FluxSetTimestepsStep()), + ("prepare_rope_inputs", FluxRoPEInputsStep()), ] - block_names = ["input", "prepare_latents", "set_timesteps"] +) + + +class FluxBeforeDenoiseStep(SequentialPipelineBlocks): + block_classes = FluxBeforeDenoiseBlocks.values() + block_names = FluxBeforeDenoiseBlocks.keys() @property def description(self): - return ( - "Before denoise step that prepare the inputs for the denoise step.\n" - + "This is a sequential pipeline blocks:\n" - + " - `FluxInputStep` is used to adjust the batch size of the model inputs\n" - + " - `FluxPrepareLatentsStep` is used to prepare the latents\n" - + " - `FluxSetTimestepsStep` is used to set the timesteps\n" - ) + return "Before denoise step that prepares the inputs for the denoise step in text-to-image generation." # before_denoise: img2img +FluxImg2ImgBeforeDenoiseBlocks = InsertableDict( + [ + ("prepare_latents", FluxPrepareLatentsStep()), + ("set_timesteps", FluxImg2ImgSetTimestepsStep()), + ("prepare_img2img_latents", FluxImg2ImgPrepareLatentsStep()), + ("prepare_rope_inputs", FluxRoPEInputsStep()), + ] +) + + class FluxImg2ImgBeforeDenoiseStep(SequentialPipelineBlocks): - block_classes = [FluxInputStep, FluxImg2ImgSetTimestepsStep, FluxImg2ImgPrepareLatentsStep] - block_names = ["input", "set_timesteps", "prepare_latents"] + block_classes = FluxImg2ImgBeforeDenoiseBlocks.values() + block_names = FluxImg2ImgBeforeDenoiseBlocks.keys() @property def description(self): - return ( - "Before denoise step that prepare the inputs for the denoise step for img2img task.\n" - + "This is a sequential pipeline blocks:\n" - + " - `FluxInputStep` is used to adjust the batch size of the model inputs\n" - + " - `FluxImg2ImgSetTimestepsStep` is used to set the timesteps\n" - + " - `FluxImg2ImgPrepareLatentsStep` is used to prepare the latents\n" - ) + return "Before denoise step that prepare the inputs for the denoise step for img2img task." # before_denoise: all task (text2img, img2img) @@ -113,7 +136,7 @@ class FluxAutoDenoiseStep(AutoPipelineBlocks): ) -# decode: all task (text2img, img2img, inpainting) +# decode: all task (text2img, img2img) class FluxAutoDecodeStep(AutoPipelineBlocks): block_classes = [FluxDecodeStep] block_names = ["non-inpaint"] @@ -124,32 +147,73 @@ class FluxAutoDecodeStep(AutoPipelineBlocks): return "Decode step that decode the denoised latents into image outputs.\n - `FluxDecodeStep`" +# inputs: text2image/img2img +FluxImg2ImgBlocks = InsertableDict( + [("text_inputs", FluxTextInputStep()), ("additional_inputs", FluxInputsDynamicStep())] +) + + +class FluxImg2ImgInputStep(SequentialPipelineBlocks): + model_name = "flux" + block_classes = FluxImg2ImgBlocks.values() + block_names = FluxImg2ImgBlocks.keys() + + @property + def description(self): + return "Input step that prepares the inputs for the img2img denoising step. It:\n" + " - make sure the text embeddings have consistent batch size as well as the additional inputs (`image_latents`).\n" + " - update height/width based `image_latents`, patchify `image_latents`." + + +class FluxImageAutoInputStep(AutoPipelineBlocks): + block_classes = [FluxImg2ImgInputStep, FluxTextInputStep] + block_names = ["img2img", "text2image"] + block_trigger_inputs = ["image_latents", None] + + @property + def description(self): + return ( + "Input step that standardize the inputs for the denoising step, e.g. make sure inputs have consistent batch size, and patchified. \n" + " This is an auto pipeline block that works for text2image/img2img tasks.\n" + + " - `FluxImg2ImgInputStep` (img2img) is used when `image_latents` is provided.\n" + + " - `FluxTextInputStep` (text2image) is used when `image_latents` are not provided.\n" + ) + + class FluxCoreDenoiseStep(SequentialPipelineBlocks): - block_classes = [FluxInputStep, FluxAutoBeforeDenoiseStep, FluxAutoDenoiseStep] + model_name = "flux" + block_classes = [FluxImageAutoInputStep, FluxAutoBeforeDenoiseStep, FluxAutoDenoiseStep] block_names = ["input", "before_denoise", "denoise"] @property def description(self): return ( "Core step that performs the denoising process. \n" - + " - `FluxInputStep` (input) standardizes the inputs for the denoising step.\n" + + " - `FluxImageAutoInputStep` (input) standardizes the inputs for the denoising step.\n" + " - `FluxAutoBeforeDenoiseStep` (before_denoise) prepares the inputs for the denoising step.\n" + " - `FluxAutoDenoiseStep` (denoise) iteratively denoises the latents.\n" - + "This step support text-to-image and image-to-image tasks for Flux:\n" + + "This step supports text-to-image and image-to-image tasks for Flux:\n" + " - for image-to-image generation, you need to provide `image_latents`\n" - + " - for text-to-image generation, all you need to provide is prompt embeddings" + + " - for text-to-image generation, all you need to provide is prompt embeddings." ) -# text2image -class FluxAutoBlocks(SequentialPipelineBlocks): - block_classes = [ - FluxTextEncoderStep, - FluxAutoVaeEncoderStep, - FluxCoreDenoiseStep, - FluxAutoDecodeStep, +# Auto blocks (text2image and img2img) +AUTO_BLOCKS = InsertableDict( + [ + ("text_encoder", FluxTextEncoderStep()), + ("image_encoder", FluxAutoVaeEncoderStep()), + ("denoise", FluxCoreDenoiseStep()), + ("decode", FluxDecodeStep()), ] - block_names = ["text_encoder", "image_encoder", "denoise", "decode"] +) + + +class FluxAutoBlocks(SequentialPipelineBlocks): + model_name = "flux" + + block_classes = AUTO_BLOCKS.values() + block_names = AUTO_BLOCKS.keys() @property def description(self): @@ -162,35 +226,28 @@ class FluxAutoBlocks(SequentialPipelineBlocks): TEXT2IMAGE_BLOCKS = InsertableDict( [ - ("text_encoder", FluxTextEncoderStep), - ("input", FluxInputStep), - ("prepare_latents", FluxPrepareLatentsStep), - ("set_timesteps", FluxSetTimestepsStep), - ("denoise", FluxDenoiseStep), - ("decode", FluxDecodeStep), + ("text_encoder", FluxTextEncoderStep()), + ("input", FluxTextInputStep()), + ("prepare_latents", FluxPrepareLatentsStep()), + ("set_timesteps", FluxSetTimestepsStep()), + ("prepare_rope_inputs", FluxRoPEInputsStep()), + ("denoise", FluxDenoiseStep()), + ("decode", FluxDecodeStep()), ] ) IMAGE2IMAGE_BLOCKS = InsertableDict( [ - ("text_encoder", FluxTextEncoderStep), - ("image_encoder", FluxVaeEncoderStep), - ("input", FluxInputStep), - ("set_timesteps", FluxImg2ImgSetTimestepsStep), - ("prepare_latents", FluxImg2ImgPrepareLatentsStep), - ("denoise", FluxDenoiseStep), - ("decode", FluxDecodeStep), + ("text_encoder", FluxTextEncoderStep()), + ("vae_encoder", FluxVaeEncoderDynamicStep()), + ("input", FluxImg2ImgInputStep()), + ("prepare_latents", FluxPrepareLatentsStep()), + ("set_timesteps", FluxImg2ImgSetTimestepsStep()), + ("prepare_img2img_latents", FluxImg2ImgPrepareLatentsStep()), + ("prepare_rope_inputs", FluxRoPEInputsStep()), + ("denoise", FluxDenoiseStep()), + ("decode", FluxDecodeStep()), ] ) -AUTO_BLOCKS = InsertableDict( - [ - ("text_encoder", FluxTextEncoderStep), - ("image_encoder", FluxAutoVaeEncoderStep), - ("denoise", FluxCoreDenoiseStep), - ("decode", FluxAutoDecodeStep), - ] -) - - ALL_BLOCKS = {"text2image": TEXT2IMAGE_BLOCKS, "img2img": IMAGE2IMAGE_BLOCKS, "auto": AUTO_BLOCKS} From 35e538d46a32e6ef588678f478437d594c32f949 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Wed, 8 Oct 2025 09:46:18 +0530 Subject: [PATCH 056/123] fix dockerfile definitions. (#12424) * fix dockerfile definitions. * python 3.10 slim. * up * up * up * up * up * revert pr_tests.yml changes * up * up * reduce python version for torch 2.1.0 --- .github/workflows/build_docker_images.yml | 1 - .github/workflows/pr_tests.yml | 1 - docker/diffusers-doc-builder/Dockerfile | 78 ++++++++----------- docker/diffusers-pytorch-cpu/Dockerfile | 67 +++++++--------- docker/diffusers-pytorch-cuda/Dockerfile | 43 +++++----- .../diffusers-pytorch-minimum-cuda/Dockerfile | 42 +++++----- .../Dockerfile | 70 ++++++++--------- 7 files changed, 133 insertions(+), 169 deletions(-) diff --git a/.github/workflows/build_docker_images.yml b/.github/workflows/build_docker_images.yml index 583853c6d6..1d7be0d6bc 100644 --- a/.github/workflows/build_docker_images.yml +++ b/.github/workflows/build_docker_images.yml @@ -72,7 +72,6 @@ jobs: image-name: - diffusers-pytorch-cpu - diffusers-pytorch-cuda - - diffusers-pytorch-cuda - diffusers-pytorch-xformers-cuda - diffusers-pytorch-minimum-cuda - diffusers-doc-builder diff --git a/.github/workflows/pr_tests.yml b/.github/workflows/pr_tests.yml index 1543b264b0..005e89f405 100644 --- a/.github/workflows/pr_tests.yml +++ b/.github/workflows/pr_tests.yml @@ -286,4 +286,3 @@ jobs: with: name: pr_main_test_reports path: reports - diff --git a/docker/diffusers-doc-builder/Dockerfile b/docker/diffusers-doc-builder/Dockerfile index 3a76b3331c..313eb8b8f8 100644 --- a/docker/diffusers-doc-builder/Dockerfile +++ b/docker/diffusers-doc-builder/Dockerfile @@ -1,56 +1,42 @@ -FROM ubuntu:20.04 +FROM python:3.10-slim +ENV PYTHONDONTWRITEBYTECODE=1 LABEL maintainer="Hugging Face" LABEL repository="diffusers" ENV DEBIAN_FRONTEND=noninteractive -RUN apt-get -y update \ - && apt-get install -y software-properties-common \ - && add-apt-repository ppa:deadsnakes/ppa +RUN apt-get -y update && apt-get install -y bash \ + build-essential \ + git \ + git-lfs \ + curl \ + ca-certificates \ + libsndfile1-dev \ + libgl1 -RUN apt install -y bash \ - build-essential \ - git \ - git-lfs \ - curl \ - ca-certificates \ - libsndfile1-dev \ - python3.10 \ - python3-pip \ - libgl1 \ - zip \ - wget \ - python3.10-venv && \ - rm -rf /var/lib/apt/lists - -# make sure to use venv -RUN python3.10 -m venv /opt/venv -ENV PATH="/opt/venv/bin:$PATH" +ENV UV_PYTHON=/usr/local/bin/python # pre-install the heavy dependencies (these can later be overridden by the deps from setup.py) -RUN python3.10 -m pip install --no-cache-dir --upgrade pip uv==0.1.11 && \ - python3.10 -m uv pip install --no-cache-dir \ - torch \ - torchvision \ - torchaudio \ - invisible_watermark \ - --extra-index-url https://download.pytorch.org/whl/cpu && \ - python3.10 -m uv pip install --no-cache-dir \ - accelerate \ - datasets \ - hf-doc-builder \ - huggingface-hub \ - Jinja2 \ - librosa \ - numpy==1.26.4 \ - scipy \ - tensorboard \ - transformers \ - matplotlib \ - setuptools==69.5.1 \ - bitsandbytes \ - torchao \ - gguf \ - optimum-quanto +RUN pip install uv +RUN uv pip install --no-cache-dir \ + torch \ + torchvision \ + torchaudio \ + --extra-index-url https://download.pytorch.org/whl/cpu + +RUN uv pip install --no-cache-dir "git+https://github.com/huggingface/diffusers.git@main#egg=diffusers[test]" + +# Extra dependencies +RUN uv pip install --no-cache-dir \ + accelerate \ + numpy==1.26.4 \ + hf_transfer \ + setuptools==69.5.1 \ + bitsandbytes \ + torchao \ + gguf \ + optimum-quanto + +RUN apt-get clean && rm -rf /var/lib/apt/lists/* && apt-get autoremove && apt-get autoclean CMD ["/bin/bash"] diff --git a/docker/diffusers-pytorch-cpu/Dockerfile b/docker/diffusers-pytorch-cpu/Dockerfile index 8d98c52598..3fc16d57e8 100644 --- a/docker/diffusers-pytorch-cpu/Dockerfile +++ b/docker/diffusers-pytorch-cpu/Dockerfile @@ -1,50 +1,37 @@ -FROM ubuntu:20.04 +FROM python:3.10-slim +ENV PYTHONDONTWRITEBYTECODE=1 LABEL maintainer="Hugging Face" LABEL repository="diffusers" ENV DEBIAN_FRONTEND=noninteractive -RUN apt-get -y update \ - && apt-get install -y software-properties-common \ - && add-apt-repository ppa:deadsnakes/ppa +RUN apt-get -y update && apt-get install -y bash \ + build-essential \ + git \ + git-lfs \ + curl \ + ca-certificates \ + libsndfile1-dev \ + libgl1 -RUN apt install -y bash \ - build-essential \ - git \ - git-lfs \ - curl \ - ca-certificates \ - libsndfile1-dev \ - python3.10 \ - python3.10-dev \ - python3-pip \ - libgl1 \ - python3.10-venv && \ - rm -rf /var/lib/apt/lists - -# make sure to use venv -RUN python3.10 -m venv /opt/venv -ENV PATH="/opt/venv/bin:$PATH" +ENV UV_PYTHON=/usr/local/bin/python # pre-install the heavy dependencies (these can later be overridden by the deps from setup.py) -RUN python3.10 -m pip install --no-cache-dir --upgrade pip uv==0.1.11 && \ - python3.10 -m uv pip install --no-cache-dir \ - torch \ - torchvision \ - torchaudio \ - invisible_watermark \ - --extra-index-url https://download.pytorch.org/whl/cpu && \ - python3.10 -m uv pip install --no-cache-dir \ - accelerate \ - datasets \ - hf-doc-builder \ - huggingface-hub \ - Jinja2 \ - librosa \ - numpy==1.26.4 \ - scipy \ - tensorboard \ - transformers matplotlib \ - hf_transfer +RUN pip install uv +RUN uv pip install --no-cache-dir \ + torch \ + torchvision \ + torchaudio \ + --extra-index-url https://download.pytorch.org/whl/cpu + +RUN uv pip install --no-cache-dir "git+https://github.com/huggingface/diffusers.git@main#egg=diffusers[test]" + +# Extra dependencies +RUN uv pip install --no-cache-dir \ + accelerate \ + numpy==1.26.4 \ + hf_transfer + +RUN apt-get clean && rm -rf /var/lib/apt/lists/* && apt-get autoremove && apt-get autoclean CMD ["/bin/bash"] diff --git a/docker/diffusers-pytorch-cuda/Dockerfile b/docker/diffusers-pytorch-cuda/Dockerfile index 695f5ed08d..1cfcda643a 100644 --- a/docker/diffusers-pytorch-cuda/Dockerfile +++ b/docker/diffusers-pytorch-cuda/Dockerfile @@ -2,11 +2,13 @@ FROM nvidia/cuda:12.1.0-runtime-ubuntu20.04 LABEL maintainer="Hugging Face" LABEL repository="diffusers" +ARG PYTHON_VERSION=3.12 ENV DEBIAN_FRONTEND=noninteractive RUN apt-get -y update \ && apt-get install -y software-properties-common \ - && add-apt-repository ppa:deadsnakes/ppa + && add-apt-repository ppa:deadsnakes/ppa && \ + apt-get update RUN apt install -y bash \ build-essential \ @@ -16,36 +18,31 @@ RUN apt install -y bash \ ca-certificates \ libsndfile1-dev \ libgl1 \ - python3.10 \ - python3.10-dev \ + python3 \ python3-pip \ - python3.10-venv && \ - rm -rf /var/lib/apt/lists + && apt-get clean \ + && rm -rf /var/lib/apt/lists/* -# make sure to use venv -RUN python3.10 -m venv /opt/venv -ENV PATH="/opt/venv/bin:$PATH" +RUN curl -LsSf https://astral.sh/uv/install.sh | sh +ENV PATH="/root/.local/bin:$PATH" +ENV VIRTUAL_ENV="/opt/venv" +ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python +RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV} +ENV PATH="$VIRTUAL_ENV/bin:$PATH" # pre-install the heavy dependencies (these can later be overridden by the deps from setup.py) -RUN python3.10 -m pip install --no-cache-dir --upgrade pip uv==0.1.11 && \ - python3.10 -m uv pip install --no-cache-dir \ +RUN uv pip install --no-cache-dir \ torch \ torchvision \ - torchaudio \ - invisible_watermark && \ - python3.10 -m pip install --no-cache-dir \ + torchaudio + +RUN uv pip install --no-cache-dir "git+https://github.com/huggingface/diffusers.git@main#egg=diffusers[test]" + +# Extra dependencies +RUN uv pip install --no-cache-dir \ accelerate \ - datasets \ - hf-doc-builder \ - huggingface-hub \ - hf_transfer \ - Jinja2 \ - librosa \ numpy==1.26.4 \ - scipy \ - tensorboard \ - transformers \ - pytorch-lightning \ + pytorch-lightning \ hf_transfer CMD ["/bin/bash"] diff --git a/docker/diffusers-pytorch-minimum-cuda/Dockerfile b/docker/diffusers-pytorch-minimum-cuda/Dockerfile index 57ca7657ac..8ca6a157da 100644 --- a/docker/diffusers-pytorch-minimum-cuda/Dockerfile +++ b/docker/diffusers-pytorch-minimum-cuda/Dockerfile @@ -2,6 +2,7 @@ FROM nvidia/cuda:12.1.0-runtime-ubuntu20.04 LABEL maintainer="Hugging Face" LABEL repository="diffusers" +ARG PYTHON_VERSION=3.10 ENV DEBIAN_FRONTEND=noninteractive ENV MINIMUM_SUPPORTED_TORCH_VERSION="2.1.0" ENV MINIMUM_SUPPORTED_TORCHVISION_VERSION="0.16.0" @@ -9,7 +10,8 @@ ENV MINIMUM_SUPPORTED_TORCHAUDIO_VERSION="2.1.0" RUN apt-get -y update \ && apt-get install -y software-properties-common \ - && add-apt-repository ppa:deadsnakes/ppa + && add-apt-repository ppa:deadsnakes/ppa && \ + apt-get update RUN apt install -y bash \ build-essential \ @@ -19,35 +21,31 @@ RUN apt install -y bash \ ca-certificates \ libsndfile1-dev \ libgl1 \ - python3.10 \ - python3.10-dev \ + python3 \ python3-pip \ - python3.10-venv && \ - rm -rf /var/lib/apt/lists + && apt-get clean \ + && rm -rf /var/lib/apt/lists/* -# make sure to use venv -RUN python3.10 -m venv /opt/venv -ENV PATH="/opt/venv/bin:$PATH" +RUN curl -LsSf https://astral.sh/uv/install.sh | sh +ENV PATH="/root/.local/bin:$PATH" +ENV VIRTUAL_ENV="/opt/venv" +ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python +RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV} +ENV PATH="$VIRTUAL_ENV/bin:$PATH" # pre-install the heavy dependencies (these can later be overridden by the deps from setup.py) -RUN python3.10 -m pip install --no-cache-dir --upgrade pip uv==0.1.11 && \ - python3.10 -m uv pip install --no-cache-dir \ +RUN uv pip install --no-cache-dir \ torch==$MINIMUM_SUPPORTED_TORCH_VERSION \ torchvision==$MINIMUM_SUPPORTED_TORCHVISION_VERSION \ - torchaudio==$MINIMUM_SUPPORTED_TORCHAUDIO_VERSION \ - invisible_watermark && \ - python3.10 -m pip install --no-cache-dir \ + torchaudio==$MINIMUM_SUPPORTED_TORCHAUDIO_VERSION + +RUN uv pip install --no-cache-dir "git+https://github.com/huggingface/diffusers.git@main#egg=diffusers[test]" + +# Extra dependencies +RUN uv pip install --no-cache-dir \ accelerate \ - datasets \ - hf-doc-builder \ - huggingface-hub \ - hf_transfer \ - Jinja2 \ - librosa \ numpy==1.26.4 \ - scipy \ - tensorboard \ - transformers \ + pytorch-lightning \ hf_transfer CMD ["/bin/bash"] diff --git a/docker/diffusers-pytorch-xformers-cuda/Dockerfile b/docker/diffusers-pytorch-xformers-cuda/Dockerfile index 1693eb2930..fe1f475504 100644 --- a/docker/diffusers-pytorch-xformers-cuda/Dockerfile +++ b/docker/diffusers-pytorch-xformers-cuda/Dockerfile @@ -2,50 +2,48 @@ FROM nvidia/cuda:12.1.0-runtime-ubuntu20.04 LABEL maintainer="Hugging Face" LABEL repository="diffusers" +ARG PYTHON_VERSION=3.12 ENV DEBIAN_FRONTEND=noninteractive RUN apt-get -y update \ && apt-get install -y software-properties-common \ - && add-apt-repository ppa:deadsnakes/ppa + && add-apt-repository ppa:deadsnakes/ppa && \ + apt-get update RUN apt install -y bash \ - build-essential \ - git \ - git-lfs \ - curl \ - ca-certificates \ - libsndfile1-dev \ - libgl1 \ - python3.10 \ - python3.10-dev \ - python3-pip \ - python3.10-venv && \ - rm -rf /var/lib/apt/lists + build-essential \ + git \ + git-lfs \ + curl \ + ca-certificates \ + libsndfile1-dev \ + libgl1 \ + python3 \ + python3-pip \ + && apt-get clean \ + && rm -rf /var/lib/apt/lists/* -# make sure to use venv -RUN python3.10 -m venv /opt/venv -ENV PATH="/opt/venv/bin:$PATH" +RUN curl -LsSf https://astral.sh/uv/install.sh | sh +ENV PATH="/root/.local/bin:$PATH" +ENV VIRTUAL_ENV="/opt/venv" +ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python +RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV} +ENV PATH="$VIRTUAL_ENV/bin:$PATH" # pre-install the heavy dependencies (these can later be overridden by the deps from setup.py) -RUN python3.10 -m pip install --no-cache-dir --upgrade pip uv==0.1.11 && \ - python3.10 -m pip install --no-cache-dir \ - torch \ - torchvision \ - torchaudio \ - invisible_watermark && \ - python3.10 -m uv pip install --no-cache-dir \ - accelerate \ - datasets \ - hf-doc-builder \ - huggingface-hub \ - hf_transfer \ - Jinja2 \ - librosa \ - numpy==1.26.4 \ - scipy \ - tensorboard \ - transformers \ - xformers \ - hf_transfer +RUN uv pip install --no-cache-dir \ + torch \ + torchvision \ + torchaudio + +RUN uv pip install --no-cache-dir "git+https://github.com/huggingface/diffusers.git@main#egg=diffusers[test]" + +# Extra dependencies +RUN uv pip install --no-cache-dir \ + accelerate \ + numpy==1.26.4 \ + pytorch-lightning \ + hf_transfer \ + xformers CMD ["/bin/bash"] From 345864eb852b528fd1f4b6ad087fa06e0470006b Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Wed, 8 Oct 2025 10:45:39 +0530 Subject: [PATCH 057/123] fix more torch.distributed imports (#12425) * up * unguard. --- src/diffusers/hooks/context_parallel.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/diffusers/hooks/context_parallel.py b/src/diffusers/hooks/context_parallel.py index 83406d4969..915fe453b9 100644 --- a/src/diffusers/hooks/context_parallel.py +++ b/src/diffusers/hooks/context_parallel.py @@ -17,7 +17,10 @@ from dataclasses import dataclass from typing import Dict, List, Type, Union import torch -import torch.distributed._functional_collectives as funcol + + +if torch.distributed.is_available(): + import torch.distributed._functional_collectives as funcol from ..models._modeling_parallel import ( ContextParallelConfig, From a519272d97f011332588e1aaa73d32952d80f3f5 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Wed, 8 Oct 2025 19:21:24 +0530 Subject: [PATCH 058/123] [ci] revisit the installations in CI. (#12450) * revisit the installations in CI. * up * up * up * empty * up * up * up --- .github/workflows/benchmark.yml | 5 +- .../workflows/mirror_community_pipeline.yml | 2 +- .github/workflows/nightly_tests.yml | 91 +++++++++---------- .github/workflows/pr_dependency_test.yml | 9 +- .github/workflows/pr_modular_tests.yml | 15 ++- .github/workflows/pr_test_fetcher.yml | 19 ++-- .github/workflows/pr_tests.yml | 48 ++++------ .github/workflows/pr_tests_gpu.yml | 42 ++++----- .../workflows/pr_torch_dependency_test.yml | 10 +- .github/workflows/push_tests.yml | 38 +++----- .github/workflows/push_tests_fast.yml | 12 +-- .github/workflows/release_tests_fast.yml | 47 ++++------ .github/workflows/run_tests_from_a_pr.yml | 5 +- docker/diffusers-doc-builder/Dockerfile | 1 + docker/diffusers-pytorch-cpu/Dockerfile | 1 + docker/diffusers-pytorch-cuda/Dockerfile | 1 + .../diffusers-pytorch-minimum-cuda/Dockerfile | 1 + .../Dockerfile | 1 + setup.py | 3 +- src/diffusers/dependency_versions_table.py | 1 + 20 files changed, 148 insertions(+), 204 deletions(-) diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index cc97e043c1..939cce9ffb 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -38,9 +38,8 @@ jobs: run: | apt update apt install -y libpq-dev postgresql-client - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test] - python -m uv pip install -r benchmarks/requirements.txt + uv pip install -e ".[quality]" + uv pip install -r benchmarks/requirements.txt - name: Environment run: | python utils/print_env.py diff --git a/.github/workflows/mirror_community_pipeline.yml b/.github/workflows/mirror_community_pipeline.yml index 9cf573312b..ab4ded9730 100644 --- a/.github/workflows/mirror_community_pipeline.yml +++ b/.github/workflows/mirror_community_pipeline.yml @@ -74,7 +74,7 @@ jobs: python-version: "3.10" - name: Install dependencies run: | - python -m pip install --upgrade pip + pip install --upgrade pip pip install --upgrade huggingface_hub # Check secret is set diff --git a/.github/workflows/nightly_tests.yml b/.github/workflows/nightly_tests.yml index 479e5503ee..c8fa3a7ad9 100644 --- a/.github/workflows/nightly_tests.yml +++ b/.github/workflows/nightly_tests.yml @@ -71,10 +71,9 @@ jobs: run: nvidia-smi - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test] - pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git - python -m uv pip install pytest-reportlog + uv pip install -e ".[quality]" + uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git + uv pip install pytest-reportlog - name: Environment run: | python utils/print_env.py @@ -84,7 +83,7 @@ jobs: # https://pytorch.org/docs/stable/notes/randomness.html#avoiding-nondeterministic-algorithms CUBLAS_WORKSPACE_CONFIG: :16:8 run: | - python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile \ + pytest -n 1 --max-worker-restart=0 --dist=loadfile \ -s -v -k "not Flax and not Onnx" \ --make-reports=tests_pipeline_${{ matrix.module }}_cuda \ --report-log=tests_pipeline_${{ matrix.module }}_cuda.log \ @@ -124,11 +123,10 @@ jobs: - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test] - python -m uv pip install peft@git+https://github.com/huggingface/peft.git - pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git - python -m uv pip install pytest-reportlog + uv pip install -e ".[quality]" + uv pip install peft@git+https://github.com/huggingface/peft.git + uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git + uv pip install pytest-reportlog - name: Environment run: python utils/print_env.py @@ -139,7 +137,7 @@ jobs: # https://pytorch.org/docs/stable/notes/randomness.html#avoiding-nondeterministic-algorithms CUBLAS_WORKSPACE_CONFIG: :16:8 run: | - python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile \ + pytest -n 1 --max-worker-restart=0 --dist=loadfile \ -s -v -k "not Flax and not Onnx" \ --make-reports=tests_torch_${{ matrix.module }}_cuda \ --report-log=tests_torch_${{ matrix.module }}_cuda.log \ @@ -152,7 +150,7 @@ jobs: # https://pytorch.org/docs/stable/notes/randomness.html#avoiding-nondeterministic-algorithms CUBLAS_WORKSPACE_CONFIG: :16:8 run: | - python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile \ + pytest -n 1 --max-worker-restart=0 --dist=loadfile \ -s -v --make-reports=examples_torch_cuda \ --report-log=examples_torch_cuda.log \ examples/ @@ -191,8 +189,7 @@ jobs: nvidia-smi - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test,training] + uv pip install -e ".[quality,training]" - name: Environment run: | python utils/print_env.py @@ -201,7 +198,7 @@ jobs: HF_TOKEN: ${{ secrets.DIFFUSERS_HF_HUB_READ_TOKEN }} RUN_COMPILE: yes run: | - python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile -s -v -k "compile" --make-reports=tests_torch_compile_cuda tests/ + pytest -n 1 --max-worker-restart=0 --dist=loadfile -s -v -k "compile" --make-reports=tests_torch_compile_cuda tests/ - name: Failure short reports if: ${{ failure() }} run: cat reports/tests_torch_compile_cuda_failures_short.txt @@ -232,11 +229,10 @@ jobs: run: nvidia-smi - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test] - python -m uv pip install peft@git+https://github.com/huggingface/peft.git - pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git - python -m uv pip install pytest-reportlog + uv pip install -e ".[quality]" + uv pip install peft@git+https://github.com/huggingface/peft.git + uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git + uv pip install pytest-reportlog - name: Environment run: | python utils/print_env.py @@ -247,7 +243,7 @@ jobs: CUBLAS_WORKSPACE_CONFIG: :16:8 BIG_GPU_MEMORY: 40 run: | - python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile \ + pytest -n 1 --max-worker-restart=0 --dist=loadfile \ -m "big_accelerator" \ --make-reports=tests_big_gpu_torch_cuda \ --report-log=tests_big_gpu_torch_cuda.log \ @@ -282,10 +278,9 @@ jobs: - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test] - python -m uv pip install peft@git+https://github.com/huggingface/peft.git - pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git + uv pip install -e ".[quality]" + uv pip install peft@git+https://github.com/huggingface/peft.git + uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git - name: Environment run: | @@ -297,7 +292,7 @@ jobs: # https://pytorch.org/docs/stable/notes/randomness.html#avoiding-nondeterministic-algorithms CUBLAS_WORKSPACE_CONFIG: :16:8 run: | - python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile \ + pytest -n 1 --max-worker-restart=0 --dist=loadfile \ -s -v -k "not Flax and not Onnx" \ --make-reports=tests_torch_minimum_version_cuda \ tests/models/test_modeling_common.py \ @@ -357,13 +352,12 @@ jobs: run: nvidia-smi - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test] - python -m uv pip install -U ${{ matrix.config.backend }} + uv pip install -e ".[quality]" + uv pip install -U ${{ matrix.config.backend }} if [ "${{ join(matrix.config.additional_deps, ' ') }}" != "" ]; then - python -m uv pip install ${{ join(matrix.config.additional_deps, ' ') }} + uv pip install ${{ join(matrix.config.additional_deps, ' ') }} fi - python -m uv pip install pytest-reportlog + uv pip install pytest-reportlog - name: Environment run: | python utils/print_env.py @@ -374,7 +368,7 @@ jobs: CUBLAS_WORKSPACE_CONFIG: :16:8 BIG_GPU_MEMORY: 40 run: | - python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile \ + pytest -n 1 --max-worker-restart=0 --dist=loadfile \ --make-reports=tests_${{ matrix.config.backend }}_torch_cuda \ --report-log=tests_${{ matrix.config.backend }}_torch_cuda.log \ tests/quantization/${{ matrix.config.test_location }} @@ -409,10 +403,9 @@ jobs: run: nvidia-smi - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test] - python -m uv pip install -U bitsandbytes optimum_quanto - python -m uv pip install pytest-reportlog + uv pip install -e ".[quality]" + uv pip install -U bitsandbytes optimum_quanto + uv pip install pytest-reportlog - name: Environment run: | python utils/print_env.py @@ -423,7 +416,7 @@ jobs: CUBLAS_WORKSPACE_CONFIG: :16:8 BIG_GPU_MEMORY: 40 run: | - python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile \ + pytest -n 1 --max-worker-restart=0 --dist=loadfile \ --make-reports=tests_pipeline_level_quant_torch_cuda \ --report-log=tests_pipeline_level_quant_torch_cuda.log \ tests/quantization/test_pipeline_level_quantization.py @@ -523,11 +516,11 @@ jobs: # - name: Install dependencies # shell: arch -arch arm64 bash {0} # run: | -# ${CONDA_RUN} python -m pip install --upgrade pip uv -# ${CONDA_RUN} python -m uv pip install -e [quality,test] -# ${CONDA_RUN} python -m uv pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu -# ${CONDA_RUN} python -m uv pip install accelerate@git+https://github.com/huggingface/accelerate -# ${CONDA_RUN} python -m uv pip install pytest-reportlog +# ${CONDA_RUN} pip install --upgrade pip uv +# ${CONDA_RUN} uv pip install -e ".[quality]" +# ${CONDA_RUN} uv pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu +# ${CONDA_RUN} uv pip install accelerate@git+https://github.com/huggingface/accelerate +# ${CONDA_RUN} uv pip install pytest-reportlog # - name: Environment # shell: arch -arch arm64 bash {0} # run: | @@ -538,7 +531,7 @@ jobs: # HF_HOME: /System/Volumes/Data/mnt/cache # HF_TOKEN: ${{ secrets.DIFFUSERS_HF_HUB_READ_TOKEN }} # run: | -# ${CONDA_RUN} python -m pytest -n 1 -s -v --make-reports=tests_torch_mps \ +# ${CONDA_RUN} pytest -n 1 -s -v --make-reports=tests_torch_mps \ # --report-log=tests_torch_mps.log \ # tests/ # - name: Failure short reports @@ -579,11 +572,11 @@ jobs: # - name: Install dependencies # shell: arch -arch arm64 bash {0} # run: | -# ${CONDA_RUN} python -m pip install --upgrade pip uv -# ${CONDA_RUN} python -m uv pip install -e [quality,test] -# ${CONDA_RUN} python -m uv pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu -# ${CONDA_RUN} python -m uv pip install accelerate@git+https://github.com/huggingface/accelerate -# ${CONDA_RUN} python -m uv pip install pytest-reportlog +# ${CONDA_RUN} pip install --upgrade pip uv +# ${CONDA_RUN} uv pip install -e ".[quality]" +# ${CONDA_RUN} uv pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu +# ${CONDA_RUN} uv pip install accelerate@git+https://github.com/huggingface/accelerate +# ${CONDA_RUN} uv pip install pytest-reportlog # - name: Environment # shell: arch -arch arm64 bash {0} # run: | @@ -594,7 +587,7 @@ jobs: # HF_HOME: /System/Volumes/Data/mnt/cache # HF_TOKEN: ${{ secrets.DIFFUSERS_HF_HUB_READ_TOKEN }} # run: | -# ${CONDA_RUN} python -m pytest -n 1 -s -v --make-reports=tests_torch_mps \ +# ${CONDA_RUN} pytest -n 1 -s -v --make-reports=tests_torch_mps \ # --report-log=tests_torch_mps.log \ # tests/ # - name: Failure short reports diff --git a/.github/workflows/pr_dependency_test.yml b/.github/workflows/pr_dependency_test.yml index d9350c09ac..b914d10761 100644 --- a/.github/workflows/pr_dependency_test.yml +++ b/.github/workflows/pr_dependency_test.yml @@ -25,11 +25,8 @@ jobs: python-version: "3.8" - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m pip install --upgrade pip uv - python -m uv pip install -e . - python -m uv pip install pytest + pip install -e . + pip install pytest - name: Check for soft dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - pytest tests/others/test_dependencies.py + pytest tests/others/test_dependencies.py diff --git a/.github/workflows/pr_modular_tests.yml b/.github/workflows/pr_modular_tests.yml index 75258771e4..83c84fbab2 100644 --- a/.github/workflows/pr_modular_tests.yml +++ b/.github/workflows/pr_modular_tests.yml @@ -42,7 +42,7 @@ jobs: python-version: "3.10" - name: Install dependencies run: | - python -m pip install --upgrade pip + pip install --upgrade pip pip install .[quality] - name: Check quality run: make quality @@ -62,7 +62,7 @@ jobs: python-version: "3.10" - name: Install dependencies run: | - python -m pip install --upgrade pip + pip install --upgrade pip pip install .[quality] - name: Check repo consistency run: | @@ -108,21 +108,18 @@ jobs: - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test] - pip uninstall transformers -y && pip uninstall huggingface_hub -y && python -m uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git - pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git --no-deps + uv pip install -e ".[quality]" + uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git + uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git --no-deps - name: Environment run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" python utils/print_env.py - name: Run fast PyTorch Pipeline CPU tests if: ${{ matrix.config.framework == 'pytorch_pipelines' }} run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m pytest -n 8 --max-worker-restart=0 --dist=loadfile \ + pytest -n 8 --max-worker-restart=0 --dist=loadfile \ -s -v -k "not Flax and not Onnx" \ --make-reports=tests_${{ matrix.config.report }} \ tests/modular_pipelines diff --git a/.github/workflows/pr_test_fetcher.yml b/.github/workflows/pr_test_fetcher.yml index b032bb8427..83b2ab4edb 100644 --- a/.github/workflows/pr_test_fetcher.yml +++ b/.github/workflows/pr_test_fetcher.yml @@ -33,8 +33,7 @@ jobs: fetch-depth: 0 - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test] + uv pip install -e ".[quality]" - name: Environment run: | python utils/print_env.py @@ -90,19 +89,16 @@ jobs: - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m pip install -e [quality,test] - python -m pip install accelerate + uv pip install -e ".[quality]" + uv pip install accelerate - name: Environment run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" python utils/print_env.py - name: Run all selected tests on CPU run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m pytest -n 2 --dist=loadfile -v --make-reports=${{ matrix.modules }}_tests_cpu ${{ fromJson(needs.setup_pr_tests.outputs.test_map)[matrix.modules] }} + pytest -n 2 --dist=loadfile -v --make-reports=${{ matrix.modules }}_tests_cpu ${{ fromJson(needs.setup_pr_tests.outputs.test_map)[matrix.modules] }} - name: Failure short reports if: ${{ failure() }} @@ -148,19 +144,16 @@ jobs: - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m pip install -e [quality,test] + pip install -e [quality] - name: Environment run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" python utils/print_env.py - name: Run Hub tests for models, schedulers, and pipelines on a staging env if: ${{ matrix.config.framework == 'hub_tests_pytorch' }} run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - HUGGINGFACE_CO_STAGING=true python -m pytest \ + HUGGINGFACE_CO_STAGING=true pytest \ -m "is_staging_test" \ --make-reports=tests_${{ matrix.config.report }} \ tests diff --git a/.github/workflows/pr_tests.yml b/.github/workflows/pr_tests.yml index 005e89f405..03205fcec6 100644 --- a/.github/workflows/pr_tests.yml +++ b/.github/workflows/pr_tests.yml @@ -38,7 +38,7 @@ jobs: python-version: "3.8" - name: Install dependencies run: | - python -m pip install --upgrade pip + pip install --upgrade pip pip install .[quality] - name: Check quality run: make quality @@ -58,7 +58,7 @@ jobs: python-version: "3.8" - name: Install dependencies run: | - python -m pip install --upgrade pip + pip install --upgrade pip pip install .[quality] - name: Check repo consistency run: | @@ -114,21 +114,18 @@ jobs: - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test] - pip uninstall transformers -y && pip uninstall huggingface_hub -y && python -m uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git - pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git --no-deps + uv pip install -e ".[quality]" + uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git + uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git --no-deps - name: Environment run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" python utils/print_env.py - name: Run fast PyTorch Pipeline CPU tests if: ${{ matrix.config.framework == 'pytorch_pipelines' }} run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m pytest -n 8 --max-worker-restart=0 --dist=loadfile \ + pytest -n 8 --max-worker-restart=0 --dist=loadfile \ -s -v -k "not Flax and not Onnx" \ --make-reports=tests_${{ matrix.config.report }} \ tests/pipelines @@ -136,8 +133,7 @@ jobs: - name: Run fast PyTorch Model Scheduler CPU tests if: ${{ matrix.config.framework == 'pytorch_models' }} run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m pytest -n 4 --max-worker-restart=0 --dist=loadfile \ + pytest -n 4 --max-worker-restart=0 --dist=loadfile \ -s -v -k "not Flax and not Onnx and not Dependency" \ --make-reports=tests_${{ matrix.config.report }} \ tests/models tests/schedulers tests/others @@ -145,9 +141,8 @@ jobs: - name: Run example PyTorch CPU tests if: ${{ matrix.config.framework == 'pytorch_examples' }} run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install peft timm - python -m pytest -n 4 --max-worker-restart=0 --dist=loadfile \ + uv pip install ".[training]" + pytest -n 4 --max-worker-restart=0 --dist=loadfile \ --make-reports=tests_${{ matrix.config.report }} \ examples @@ -195,19 +190,16 @@ jobs: - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test] + uv pip install -e ".[quality]" - name: Environment run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" python utils/print_env.py - name: Run Hub tests for models, schedulers, and pipelines on a staging env if: ${{ matrix.config.framework == 'hub_tests_pytorch' }} run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - HUGGINGFACE_CO_STAGING=true python -m pytest \ + HUGGINGFACE_CO_STAGING=true pytest \ -m "is_staging_test" \ --make-reports=tests_${{ matrix.config.report }} \ tests @@ -249,27 +241,24 @@ jobs: - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test] + uv pip install -e ".[quality]" # TODO (sayakpaul, DN6): revisit `--no-deps` - python -m pip install -U peft@git+https://github.com/huggingface/peft.git --no-deps - python -m uv pip install -U tokenizers - pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git --no-deps - pip uninstall transformers -y && pip uninstall huggingface_hub -y && python -m uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git + uv pip install -U peft@git+https://github.com/huggingface/peft.git --no-deps + uv pip install -U tokenizers + uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git --no-deps + uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git - name: Environment run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" python utils/print_env.py - name: Run fast PyTorch LoRA tests with PEFT run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m pytest -n 4 --max-worker-restart=0 --dist=loadfile \ + pytest -n 4 --max-worker-restart=0 --dist=loadfile \ -s -v \ --make-reports=tests_peft_main \ tests/lora/ - python -m pytest -n 4 --max-worker-restart=0 --dist=loadfile \ + pytest -n 4 --max-worker-restart=0 --dist=loadfile \ -s -v \ --make-reports=tests_models_lora_peft_main \ tests/models/ -k "lora" @@ -286,3 +275,4 @@ jobs: with: name: pr_main_test_reports path: reports + diff --git a/.github/workflows/pr_tests_gpu.yml b/.github/workflows/pr_tests_gpu.yml index 89b6abe20d..900a53da94 100644 --- a/.github/workflows/pr_tests_gpu.yml +++ b/.github/workflows/pr_tests_gpu.yml @@ -39,7 +39,7 @@ jobs: python-version: "3.8" - name: Install dependencies run: | - python -m pip install --upgrade pip + pip install --upgrade pip pip install .[quality] - name: Check quality run: make quality @@ -59,7 +59,7 @@ jobs: python-version: "3.8" - name: Install dependencies run: | - python -m pip install --upgrade pip + pip install --upgrade pip pip install .[quality] - name: Check repo consistency run: | @@ -88,8 +88,7 @@ jobs: fetch-depth: 2 - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test] + uv pip install -e ".[quality]" - name: Environment run: | python utils/print_env.py @@ -130,10 +129,9 @@ jobs: nvidia-smi - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test] - pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git - pip uninstall transformers -y && pip uninstall huggingface_hub -y && python -m uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git + uv pip install -e ".[quality]" + uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git + uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git - name: Environment run: | @@ -152,13 +150,13 @@ jobs: CUBLAS_WORKSPACE_CONFIG: :16:8 run: | if [ "${{ matrix.module }}" = "ip_adapters" ]; then - python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile \ + pytest -n 1 --max-worker-restart=0 --dist=loadfile \ -s -v -k "not Flax and not Onnx" \ --make-reports=tests_pipeline_${{ matrix.module }}_cuda \ tests/pipelines/${{ matrix.module }} else pattern=$(cat ${{ steps.extract_tests.outputs.pattern_file }}) - python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile \ + pytest -n 1 --max-worker-restart=0 --dist=loadfile \ -s -v -k "not Flax and not Onnx and $pattern" \ --make-reports=tests_pipeline_${{ matrix.module }}_cuda \ tests/pipelines/${{ matrix.module }} @@ -200,11 +198,10 @@ jobs: - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test] - python -m uv pip install peft@git+https://github.com/huggingface/peft.git - pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git - pip uninstall transformers -y && pip uninstall huggingface_hub -y && python -m uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git + uv pip install -e ".[quality]" + uv pip install peft@git+https://github.com/huggingface/peft.git + uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git + uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git - name: Environment run: | @@ -225,10 +222,10 @@ jobs: run: | pattern=$(cat ${{ steps.extract_tests.outputs.pattern_file }}) if [ -z "$pattern" ]; then - python -m pytest -n 1 -sv --max-worker-restart=0 --dist=loadfile -k "not Flax and not Onnx" tests/${{ matrix.module }} \ + pytest -n 1 -sv --max-worker-restart=0 --dist=loadfile -k "not Flax and not Onnx" tests/${{ matrix.module }} \ --make-reports=tests_torch_cuda_${{ matrix.module }} else - python -m pytest -n 1 -sv --max-worker-restart=0 --dist=loadfile -k "not Flax and not Onnx and $pattern" tests/${{ matrix.module }} \ + pytest -n 1 -sv --max-worker-restart=0 --dist=loadfile -k "not Flax and not Onnx and $pattern" tests/${{ matrix.module }} \ --make-reports=tests_torch_cuda_${{ matrix.module }} fi @@ -265,22 +262,19 @@ jobs: nvidia-smi - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - pip uninstall transformers -y && pip uninstall huggingface_hub -y && python -m uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git - python -m uv pip install -e [quality,test,training] + uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git + uv pip install -e ".[quality,training]" - name: Environment run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" python utils/print_env.py - name: Run example tests on GPU env: HF_TOKEN: ${{ secrets.DIFFUSERS_HF_HUB_READ_TOKEN }} run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install timm - python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile -s -v --make-reports=examples_torch_cuda examples/ + uv pip install ".[training]" + pytest -n 1 --max-worker-restart=0 --dist=loadfile -s -v --make-reports=examples_torch_cuda examples/ - name: Failure short reports if: ${{ failure() }} diff --git a/.github/workflows/pr_torch_dependency_test.yml b/.github/workflows/pr_torch_dependency_test.yml index c39d5eca2d..4b6160ff71 100644 --- a/.github/workflows/pr_torch_dependency_test.yml +++ b/.github/workflows/pr_torch_dependency_test.yml @@ -25,12 +25,8 @@ jobs: python-version: "3.8" - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m pip install --upgrade pip uv - python -m uv pip install -e . - python -m uv pip install torch torchvision torchaudio - python -m uv pip install pytest + pip install -e . + pip install torch torchvision torchaudio pytest - name: Check for soft dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - pytest tests/others/test_dependencies.py + pytest tests/others/test_dependencies.py diff --git a/.github/workflows/push_tests.yml b/.github/workflows/push_tests.yml index 6896e0145c..a1f5e12153 100644 --- a/.github/workflows/push_tests.yml +++ b/.github/workflows/push_tests.yml @@ -34,8 +34,7 @@ jobs: fetch-depth: 2 - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test] + uv pip install -e ".[quality]" - name: Environment run: | python utils/print_env.py @@ -75,9 +74,8 @@ jobs: nvidia-smi - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test] - pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git + uv pip install -e ".[quality]" + uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git - name: Environment run: | python utils/print_env.py @@ -87,7 +85,7 @@ jobs: # https://pytorch.org/docs/stable/notes/randomness.html#avoiding-nondeterministic-algorithms CUBLAS_WORKSPACE_CONFIG: :16:8 run: | - python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile \ + pytest -n 1 --max-worker-restart=0 --dist=loadfile \ -s -v -k "not Flax and not Onnx" \ --make-reports=tests_pipeline_${{ matrix.module }}_cuda \ tests/pipelines/${{ matrix.module }} @@ -126,10 +124,9 @@ jobs: - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test] - python -m uv pip install peft@git+https://github.com/huggingface/peft.git - pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git + uv pip install -e ".[quality]" + uv pip install peft@git+https://github.com/huggingface/peft.git + uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git - name: Environment run: | @@ -141,7 +138,7 @@ jobs: # https://pytorch.org/docs/stable/notes/randomness.html#avoiding-nondeterministic-algorithms CUBLAS_WORKSPACE_CONFIG: :16:8 run: | - python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile \ + pytest -n 1 --max-worker-restart=0 --dist=loadfile \ -s -v -k "not Flax and not Onnx" \ --make-reports=tests_torch_cuda_${{ matrix.module }} \ tests/${{ matrix.module }} @@ -180,8 +177,7 @@ jobs: nvidia-smi - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test,training] + uv pip install -e ".[quality,training]" - name: Environment run: | python utils/print_env.py @@ -190,7 +186,7 @@ jobs: HF_TOKEN: ${{ secrets.DIFFUSERS_HF_HUB_READ_TOKEN }} RUN_COMPILE: yes run: | - python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile -s -v -k "compile" --make-reports=tests_torch_compile_cuda tests/ + pytest -n 1 --max-worker-restart=0 --dist=loadfile -s -v -k "compile" --make-reports=tests_torch_compile_cuda tests/ - name: Failure short reports if: ${{ failure() }} run: cat reports/tests_torch_compile_cuda_failures_short.txt @@ -223,8 +219,7 @@ jobs: nvidia-smi - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test,training] + uv pip install -e ".[quality,training]" - name: Environment run: | python utils/print_env.py @@ -232,7 +227,7 @@ jobs: env: HF_TOKEN: ${{ secrets.DIFFUSERS_HF_HUB_READ_TOKEN }} run: | - python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile -s -v -k "xformers" --make-reports=tests_torch_xformers_cuda tests/ + pytest -n 1 --max-worker-restart=0 --dist=loadfile -s -v -k "xformers" --make-reports=tests_torch_xformers_cuda tests/ - name: Failure short reports if: ${{ failure() }} run: cat reports/tests_torch_xformers_cuda_failures_short.txt @@ -264,21 +259,18 @@ jobs: nvidia-smi - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test,training] + uv pip install -e ".[quality,training]" - name: Environment run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" python utils/print_env.py - name: Run example tests on GPU env: HF_TOKEN: ${{ secrets.DIFFUSERS_HF_HUB_READ_TOKEN }} run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install timm - python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile -s -v --make-reports=examples_torch_cuda examples/ + uv pip install ".[training]" + pytest -n 1 --max-worker-restart=0 --dist=loadfile -s -v --make-reports=examples_torch_cuda examples/ - name: Failure short reports if: ${{ failure() }} diff --git a/.github/workflows/push_tests_fast.yml b/.github/workflows/push_tests_fast.yml index e274cb0218..6b1dd0b2d0 100644 --- a/.github/workflows/push_tests_fast.yml +++ b/.github/workflows/push_tests_fast.yml @@ -60,19 +60,16 @@ jobs: - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test] + uv pip install -e ".[quality]" - name: Environment run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" python utils/print_env.py - name: Run fast PyTorch CPU tests if: ${{ matrix.config.framework == 'pytorch' }} run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m pytest -n 4 --max-worker-restart=0 --dist=loadfile \ + pytest -n 4 --max-worker-restart=0 --dist=loadfile \ -s -v -k "not Flax and not Onnx" \ --make-reports=tests_${{ matrix.config.report }} \ tests/ @@ -80,9 +77,8 @@ jobs: - name: Run example PyTorch CPU tests if: ${{ matrix.config.framework == 'pytorch_examples' }} run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install peft timm - python -m pytest -n 4 --max-worker-restart=0 --dist=loadfile \ + uv pip install ".[training]" + pytest -n 4 --max-worker-restart=0 --dist=loadfile \ --make-reports=tests_${{ matrix.config.report }} \ examples diff --git a/.github/workflows/release_tests_fast.yml b/.github/workflows/release_tests_fast.yml index 81a34f7a46..808818bead 100644 --- a/.github/workflows/release_tests_fast.yml +++ b/.github/workflows/release_tests_fast.yml @@ -32,8 +32,7 @@ jobs: fetch-depth: 2 - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test] + uv pip install -e ".[quality]" - name: Environment run: | python utils/print_env.py @@ -73,9 +72,8 @@ jobs: nvidia-smi - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test] - pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git + uv pip install -e ".[quality]" + uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git - name: Environment run: | python utils/print_env.py @@ -85,7 +83,7 @@ jobs: # https://pytorch.org/docs/stable/notes/randomness.html#avoiding-nondeterministic-algorithms CUBLAS_WORKSPACE_CONFIG: :16:8 run: | - python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile \ + pytest -n 1 --max-worker-restart=0 --dist=loadfile \ -s -v -k "not Flax and not Onnx" \ --make-reports=tests_pipeline_${{ matrix.module }}_cuda \ tests/pipelines/${{ matrix.module }} @@ -124,10 +122,9 @@ jobs: - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test] - python -m uv pip install peft@git+https://github.com/huggingface/peft.git - pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git + uv pip install -e ".[quality]" + uv pip install peft@git+https://github.com/huggingface/peft.git + uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git - name: Environment run: | @@ -139,7 +136,7 @@ jobs: # https://pytorch.org/docs/stable/notes/randomness.html#avoiding-nondeterministic-algorithms CUBLAS_WORKSPACE_CONFIG: :16:8 run: | - python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile \ + pytest -n 1 --max-worker-restart=0 --dist=loadfile \ -s -v -k "not Flax and not Onnx" \ --make-reports=tests_torch_${{ matrix.module }}_cuda \ tests/${{ matrix.module }} @@ -175,10 +172,9 @@ jobs: - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test] - python -m uv pip install peft@git+https://github.com/huggingface/peft.git - pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git + uv pip install -e ".[quality]" + uv pip install peft@git+https://github.com/huggingface/peft.git + uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git - name: Environment run: | @@ -190,7 +186,7 @@ jobs: # https://pytorch.org/docs/stable/notes/randomness.html#avoiding-nondeterministic-algorithms CUBLAS_WORKSPACE_CONFIG: :16:8 run: | - python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile \ + pytest -n 1 --max-worker-restart=0 --dist=loadfile \ -s -v -k "not Flax and not Onnx" \ --make-reports=tests_torch_minimum_cuda \ tests/models/test_modeling_common.py \ @@ -235,8 +231,7 @@ jobs: nvidia-smi - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test,training] + uv pip install -e ".[quality,training]" - name: Environment run: | python utils/print_env.py @@ -245,7 +240,7 @@ jobs: HF_TOKEN: ${{ secrets.DIFFUSERS_HF_HUB_READ_TOKEN }} RUN_COMPILE: yes run: | - python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile -s -v -k "compile" --make-reports=tests_torch_compile_cuda tests/ + pytest -n 1 --max-worker-restart=0 --dist=loadfile -s -v -k "compile" --make-reports=tests_torch_compile_cuda tests/ - name: Failure short reports if: ${{ failure() }} run: cat reports/tests_torch_compile_cuda_failures_short.txt @@ -278,8 +273,7 @@ jobs: nvidia-smi - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test,training] + uv pip install -e ".[quality,training]" - name: Environment run: | python utils/print_env.py @@ -287,7 +281,7 @@ jobs: env: HF_TOKEN: ${{ secrets.DIFFUSERS_HF_HUB_READ_TOKEN }} run: | - python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile -s -v -k "xformers" --make-reports=tests_torch_xformers_cuda tests/ + pytest -n 1 --max-worker-restart=0 --dist=loadfile -s -v -k "xformers" --make-reports=tests_torch_xformers_cuda tests/ - name: Failure short reports if: ${{ failure() }} run: cat reports/tests_torch_xformers_cuda_failures_short.txt @@ -321,21 +315,18 @@ jobs: - name: Install dependencies run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test,training] + uv pip install -e ".[quality,training]" - name: Environment run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" python utils/print_env.py - name: Run example tests on GPU env: HF_TOKEN: ${{ secrets.DIFFUSERS_HF_HUB_READ_TOKEN }} run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install timm - python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile -s -v --make-reports=examples_torch_cuda examples/ + uv pip install ".[training]" + pytest -n 1 --max-worker-restart=0 --dist=loadfile -s -v --make-reports=examples_torch_cuda examples/ - name: Failure short reports if: ${{ failure() }} diff --git a/.github/workflows/run_tests_from_a_pr.yml b/.github/workflows/run_tests_from_a_pr.yml index c8eee8dbbc..fa8c579dd7 100644 --- a/.github/workflows/run_tests_from_a_pr.yml +++ b/.github/workflows/run_tests_from_a_pr.yml @@ -63,9 +63,8 @@ jobs: - name: Install pytest run: | - python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m uv pip install -e [quality,test] - python -m uv pip install peft + uv pip install -e ".[quality]" + uv pip install peft - name: Run tests env: diff --git a/docker/diffusers-doc-builder/Dockerfile b/docker/diffusers-doc-builder/Dockerfile index 313eb8b8f8..4ea2829309 100644 --- a/docker/diffusers-doc-builder/Dockerfile +++ b/docker/diffusers-doc-builder/Dockerfile @@ -11,6 +11,7 @@ RUN apt-get -y update && apt-get install -y bash \ git-lfs \ curl \ ca-certificates \ + libglib2.0-0 \ libsndfile1-dev \ libgl1 diff --git a/docker/diffusers-pytorch-cpu/Dockerfile b/docker/diffusers-pytorch-cpu/Dockerfile index 3fc16d57e8..dc5d0fc71c 100644 --- a/docker/diffusers-pytorch-cpu/Dockerfile +++ b/docker/diffusers-pytorch-cpu/Dockerfile @@ -11,6 +11,7 @@ RUN apt-get -y update && apt-get install -y bash \ git-lfs \ curl \ ca-certificates \ + libglib2.0-0 \ libsndfile1-dev \ libgl1 diff --git a/docker/diffusers-pytorch-cuda/Dockerfile b/docker/diffusers-pytorch-cuda/Dockerfile index 1cfcda643a..2bdfd409b4 100644 --- a/docker/diffusers-pytorch-cuda/Dockerfile +++ b/docker/diffusers-pytorch-cuda/Dockerfile @@ -16,6 +16,7 @@ RUN apt install -y bash \ git-lfs \ curl \ ca-certificates \ + libglib2.0-0 \ libsndfile1-dev \ libgl1 \ python3 \ diff --git a/docker/diffusers-pytorch-minimum-cuda/Dockerfile b/docker/diffusers-pytorch-minimum-cuda/Dockerfile index 8ca6a157da..a2ce193f68 100644 --- a/docker/diffusers-pytorch-minimum-cuda/Dockerfile +++ b/docker/diffusers-pytorch-minimum-cuda/Dockerfile @@ -19,6 +19,7 @@ RUN apt install -y bash \ git-lfs \ curl \ ca-certificates \ + libglib2.0-0 \ libsndfile1-dev \ libgl1 \ python3 \ diff --git a/docker/diffusers-pytorch-xformers-cuda/Dockerfile b/docker/diffusers-pytorch-xformers-cuda/Dockerfile index fe1f475504..1ea258bdb7 100644 --- a/docker/diffusers-pytorch-xformers-cuda/Dockerfile +++ b/docker/diffusers-pytorch-xformers-cuda/Dockerfile @@ -16,6 +16,7 @@ RUN apt install -y bash \ git-lfs \ curl \ ca-certificates \ + libglib2.0-0 \ libsndfile1-dev \ libgl1 \ python3 \ diff --git a/setup.py b/setup.py index 372a568595..8d346ddfec 100644 --- a/setup.py +++ b/setup.py @@ -145,6 +145,7 @@ _deps = [ "black", "phonemizer", "opencv-python", + "timm", ] # this is a lookup table with items like: @@ -218,7 +219,7 @@ class DepsTableUpdateCommand(Command): extras = {} extras["quality"] = deps_list("urllib3", "isort", "ruff", "hf-doc-builder") extras["docs"] = deps_list("hf-doc-builder") -extras["training"] = deps_list("accelerate", "datasets", "protobuf", "tensorboard", "Jinja2", "peft") +extras["training"] = deps_list("accelerate", "datasets", "protobuf", "tensorboard", "Jinja2", "peft", "timm") extras["test"] = deps_list( "compel", "GitPython", diff --git a/src/diffusers/dependency_versions_table.py b/src/diffusers/dependency_versions_table.py index bfc4e9818b..6e5ac630ab 100644 --- a/src/diffusers/dependency_versions_table.py +++ b/src/diffusers/dependency_versions_table.py @@ -52,4 +52,5 @@ deps = { "black": "black", "phonemizer": "phonemizer", "opencv-python": "opencv-python", + "timm": "timm", } From a9df12ab451be65060f6f69e201e3a6cb2ee15e8 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 9 Oct 2025 15:25:03 +0530 Subject: [PATCH 059/123] Update Dockerfile to include zip wget for doc-builder (#12451) --- docker/diffusers-doc-builder/Dockerfile | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docker/diffusers-doc-builder/Dockerfile b/docker/diffusers-doc-builder/Dockerfile index 4ea2829309..8453ef4e6c 100644 --- a/docker/diffusers-doc-builder/Dockerfile +++ b/docker/diffusers-doc-builder/Dockerfile @@ -13,7 +13,9 @@ RUN apt-get -y update && apt-get install -y bash \ ca-certificates \ libglib2.0-0 \ libsndfile1-dev \ - libgl1 + libgl1 \ + zip \ + wget ENV UV_PYTHON=/usr/local/bin/python From 693d8a3a52252153dc0f1503ea87db89d2364693 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Fri, 10 Oct 2025 18:10:17 +0530 Subject: [PATCH 060/123] [modular] i2i and t2i support for kontext modular (#12454) * up * get ready * fix import * up * up --- src/diffusers/__init__.py | 4 + src/diffusers/modular_pipelines/__init__.py | 9 +- .../modular_pipelines/flux/__init__.py | 15 +- .../modular_pipelines/flux/before_denoise.py | 84 ++++++- .../modular_pipelines/flux/denoise.py | 107 +++++++++ .../modular_pipelines/flux/encoders.py | 107 +++++++-- .../modular_pipelines/flux/inputs.py | 123 ++++++++++ .../modular_pipelines/flux/modular_blocks.py | 215 +++++++++++++++++- .../flux/modular_pipeline.py | 10 + .../modular_pipelines/modular_pipeline.py | 1 + .../dummy_torch_and_transformers_objects.py | 30 +++ 11 files changed, 659 insertions(+), 46 deletions(-) diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 686e8d99da..95d559ff75 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -386,6 +386,8 @@ else: _import_structure["modular_pipelines"].extend( [ "FluxAutoBlocks", + "FluxKontextAutoBlocks", + "FluxKontextModularPipeline", "FluxModularPipeline", "QwenImageAutoBlocks", "QwenImageEditAutoBlocks", @@ -1050,6 +1052,8 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: else: from .modular_pipelines import ( FluxAutoBlocks, + FluxKontextAutoBlocks, + FluxKontextModularPipeline, FluxModularPipeline, QwenImageAutoBlocks, QwenImageEditAutoBlocks, diff --git a/src/diffusers/modular_pipelines/__init__.py b/src/diffusers/modular_pipelines/__init__.py index 2e590594af..86ed735134 100644 --- a/src/diffusers/modular_pipelines/__init__.py +++ b/src/diffusers/modular_pipelines/__init__.py @@ -46,7 +46,12 @@ else: ] _import_structure["stable_diffusion_xl"] = ["StableDiffusionXLAutoBlocks", "StableDiffusionXLModularPipeline"] _import_structure["wan"] = ["WanAutoBlocks", "WanModularPipeline"] - _import_structure["flux"] = ["FluxAutoBlocks", "FluxModularPipeline"] + _import_structure["flux"] = [ + "FluxAutoBlocks", + "FluxModularPipeline", + "FluxKontextAutoBlocks", + "FluxKontextModularPipeline", + ] _import_structure["qwenimage"] = [ "QwenImageAutoBlocks", "QwenImageModularPipeline", @@ -65,7 +70,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: from ..utils.dummy_pt_objects import * # noqa F403 else: from .components_manager import ComponentsManager - from .flux import FluxAutoBlocks, FluxModularPipeline + from .flux import FluxAutoBlocks, FluxKontextAutoBlocks, FluxKontextModularPipeline, FluxModularPipeline from .modular_pipeline import ( AutoPipelineBlocks, BlockState, diff --git a/src/diffusers/modular_pipelines/flux/__init__.py b/src/diffusers/modular_pipelines/flux/__init__.py index 2891edf790..ec00986611 100644 --- a/src/diffusers/modular_pipelines/flux/__init__.py +++ b/src/diffusers/modular_pipelines/flux/__init__.py @@ -25,14 +25,18 @@ else: _import_structure["modular_blocks"] = [ "ALL_BLOCKS", "AUTO_BLOCKS", + "AUTO_BLOCKS_KONTEXT", + "FLUX_KONTEXT_BLOCKS", "TEXT2IMAGE_BLOCKS", "FluxAutoBeforeDenoiseStep", "FluxAutoBlocks", - "FluxAutoBlocks", "FluxAutoDecodeStep", "FluxAutoDenoiseStep", + "FluxKontextAutoBlocks", + "FluxKontextAutoDenoiseStep", + "FluxKontextBeforeDenoiseStep", ] - _import_structure["modular_pipeline"] = ["FluxModularPipeline"] + _import_structure["modular_pipeline"] = ["FluxKontextModularPipeline", "FluxModularPipeline"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: try: @@ -45,13 +49,18 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: from .modular_blocks import ( ALL_BLOCKS, AUTO_BLOCKS, + AUTO_BLOCKS_KONTEXT, + FLUX_KONTEXT_BLOCKS, TEXT2IMAGE_BLOCKS, FluxAutoBeforeDenoiseStep, FluxAutoBlocks, FluxAutoDecodeStep, FluxAutoDenoiseStep, + FluxKontextAutoBlocks, + FluxKontextAutoDenoiseStep, + FluxKontextBeforeDenoiseStep, ) - from .modular_pipeline import FluxModularPipeline + from .modular_pipeline import FluxKontextModularPipeline, FluxModularPipeline else: import sys diff --git a/src/diffusers/modular_pipelines/flux/before_denoise.py b/src/diffusers/modular_pipelines/flux/before_denoise.py index 5f3193af0e..c098b7d4f1 100644 --- a/src/diffusers/modular_pipelines/flux/before_denoise.py +++ b/src/diffusers/modular_pipelines/flux/before_denoise.py @@ -118,15 +118,6 @@ def retrieve_latents( raise AttributeError("Could not access latents of provided encoder_output") -# TODO: align this with Qwen patchifier -def _pack_latents(latents, batch_size, num_channels_latents, height, width): - latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2) - latents = latents.permute(0, 2, 4, 1, 3, 5) - latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4) - - return latents - - def _get_initial_timesteps_and_optionals( transformer, scheduler, @@ -398,16 +389,15 @@ class FluxPrepareLatentsStep(ModularPipelineBlocks): f" size of {batch_size}. Make sure the batch size matches the length of the generators." ) - # TODO: move packing latents code to a patchifier + # TODO: move packing latents code to a patchifier similar to Qwen latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - latents = _pack_latents(latents, batch_size, num_channels_latents, height, width) + latents = FluxPipeline._pack_latents(latents, batch_size, num_channels_latents, height, width) return latents @torch.no_grad() def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) - block_state.height = block_state.height or components.default_height block_state.width = block_state.width or components.default_width block_state.device = components._execution_device @@ -557,3 +547,73 @@ class FluxRoPEInputsStep(ModularPipelineBlocks): self.set_block_state(state, block_state) return components, state + + +class FluxKontextRoPEInputsStep(ModularPipelineBlocks): + model_name = "flux-kontext" + + @property + def description(self) -> str: + return "Step that prepares the RoPE inputs for the denoising process of Flux Kontext. Should be placed after text encoder and latent preparation steps." + + @property + def inputs(self) -> List[InputParam]: + return [ + InputParam(name="image_height"), + InputParam(name="image_width"), + InputParam(name="height"), + InputParam(name="width"), + InputParam(name="prompt_embeds"), + ] + + @property + def intermediate_outputs(self) -> List[OutputParam]: + return [ + OutputParam( + name="txt_ids", + kwargs_type="denoiser_input_fields", + type_hint=List[int], + description="The sequence lengths of the prompt embeds, used for RoPE calculation.", + ), + OutputParam( + name="img_ids", + kwargs_type="denoiser_input_fields", + type_hint=List[int], + description="The sequence lengths of the image latents, used for RoPE calculation.", + ), + ] + + def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + prompt_embeds = block_state.prompt_embeds + device, dtype = prompt_embeds.device, prompt_embeds.dtype + block_state.txt_ids = torch.zeros(prompt_embeds.shape[1], 3).to( + device=prompt_embeds.device, dtype=prompt_embeds.dtype + ) + + img_ids = None + if ( + getattr(block_state, "image_height", None) is not None + and getattr(block_state, "image_width", None) is not None + ): + image_latent_height = 2 * (int(block_state.image_height) // (components.vae_scale_factor * 2)) + image_latent_width = 2 * (int(block_state.width) // (components.vae_scale_factor * 2)) + img_ids = FluxPipeline._prepare_latent_image_ids( + None, image_latent_height // 2, image_latent_width // 2, device, dtype + ) + # image ids are the same as latent ids with the first dimension set to 1 instead of 0 + img_ids[..., 0] = 1 + + height = 2 * (int(block_state.height) // (components.vae_scale_factor * 2)) + width = 2 * (int(block_state.width) // (components.vae_scale_factor * 2)) + latent_ids = FluxPipeline._prepare_latent_image_ids(None, height // 2, width // 2, device, dtype) + + if img_ids is not None: + latent_ids = torch.cat([latent_ids, img_ids], dim=0) + + block_state.img_ids = latent_ids + + self.set_block_state(state, block_state) + + return components, state diff --git a/src/diffusers/modular_pipelines/flux/denoise.py b/src/diffusers/modular_pipelines/flux/denoise.py index e482c198e8..b1796bb63c 100644 --- a/src/diffusers/modular_pipelines/flux/denoise.py +++ b/src/diffusers/modular_pipelines/flux/denoise.py @@ -109,6 +109,96 @@ class FluxLoopDenoiser(ModularPipelineBlocks): return components, block_state +class FluxKontextLoopDenoiser(ModularPipelineBlocks): + model_name = "flux-kontext" + + @property + def expected_components(self) -> List[ComponentSpec]: + return [ComponentSpec("transformer", FluxTransformer2DModel)] + + @property + def description(self) -> str: + return ( + "Step within the denoising loop that denoise the latents for Flux Kontext. " + "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` " + "object (e.g. `FluxDenoiseLoopWrapper`)" + ) + + @property + def inputs(self) -> List[Tuple[str, Any]]: + return [ + InputParam("joint_attention_kwargs"), + InputParam( + "latents", + required=True, + type_hint=torch.Tensor, + description="The initial latents to use for the denoising process. Can be generated in prepare_latent step.", + ), + InputParam( + "image_latents", + type_hint=torch.Tensor, + description="Image latents to use for the denoising process. Can be generated in prepare_latent step.", + ), + InputParam( + "guidance", + required=True, + type_hint=torch.Tensor, + description="Guidance scale as a tensor", + ), + InputParam( + "prompt_embeds", + required=True, + type_hint=torch.Tensor, + description="Prompt embeddings", + ), + InputParam( + "pooled_prompt_embeds", + required=True, + type_hint=torch.Tensor, + description="Pooled prompt embeddings", + ), + InputParam( + "txt_ids", + required=True, + type_hint=torch.Tensor, + description="IDs computed from text sequence needed for RoPE", + ), + InputParam( + "img_ids", + required=True, + type_hint=torch.Tensor, + description="IDs computed from latent sequence needed for RoPE", + ), + ] + + @torch.no_grad() + def __call__( + self, components: FluxModularPipeline, block_state: BlockState, i: int, t: torch.Tensor + ) -> PipelineState: + latents = block_state.latents + latent_model_input = latents + image_latents = block_state.image_latents + if image_latents is not None: + latent_model_input = torch.cat([latent_model_input, image_latents], dim=1) + + timestep = t.expand(latents.shape[0]).to(latents.dtype) + noise_pred = components.transformer( + hidden_states=latent_model_input, + timestep=timestep / 1000, + guidance=block_state.guidance, + encoder_hidden_states=block_state.prompt_embeds, + pooled_projections=block_state.pooled_prompt_embeds, + joint_attention_kwargs=block_state.joint_attention_kwargs, + txt_ids=block_state.txt_ids, + img_ids=block_state.img_ids, + return_dict=False, + )[0] + noise_pred = noise_pred[:, : latents.size(1)] + block_state.noise_pred = noise_pred + + return components, block_state + + class FluxLoopAfterDenoiser(ModularPipelineBlocks): model_name = "flux" @@ -221,3 +311,20 @@ class FluxDenoiseStep(FluxDenoiseLoopWrapper): " - `FluxLoopAfterDenoiser`\n" "This block supports both text2image and img2img tasks." ) + + +class FluxKontextDenoiseStep(FluxDenoiseLoopWrapper): + model_name = "flux-kontext" + block_classes = [FluxKontextLoopDenoiser, FluxLoopAfterDenoiser] + block_names = ["denoiser", "after_denoiser"] + + @property + def description(self) -> str: + return ( + "Denoise step that iteratively denoise the latents. \n" + "Its loop logic is defined in `FluxDenoiseLoopWrapper.__call__` method \n" + "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n" + " - `FluxKontextLoopDenoiser`\n" + " - `FluxLoopAfterDenoiser`\n" + "This block supports both text2image and img2img tasks." + ) diff --git a/src/diffusers/modular_pipelines/flux/encoders.py b/src/diffusers/modular_pipelines/flux/encoders.py index 6368086cbb..b71962bd93 100644 --- a/src/diffusers/modular_pipelines/flux/encoders.py +++ b/src/diffusers/modular_pipelines/flux/encoders.py @@ -20,7 +20,7 @@ import torch from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast from ...configuration_utils import FrozenDict -from ...image_processor import VaeImageProcessor +from ...image_processor import VaeImageProcessor, is_valid_image, is_valid_image_imagelist from ...loaders import FluxLoraLoaderMixin, TextualInversionLoaderMixin from ...models import AutoencoderKL from ...utils import USE_PEFT_BACKEND, is_ftfy_available, logging, scale_lora_layers, unscale_lora_layers @@ -83,11 +83,11 @@ def encode_vae_image(vae: AutoencoderKL, image: torch.Tensor, generator: torch.G class FluxProcessImagesInputStep(ModularPipelineBlocks): - model_name = "Flux" + model_name = "flux" @property def description(self) -> str: - return "Image Preprocess step. Resizing is needed in Flux Kontext (will be implemented later.)" + return "Image Preprocess step." @property def expected_components(self) -> List[ComponentSpec]: @@ -106,9 +106,7 @@ class FluxProcessImagesInputStep(ModularPipelineBlocks): @property def intermediate_outputs(self) -> List[OutputParam]: - return [ - OutputParam(name="processed_image"), - ] + return [OutputParam(name="processed_image")] @staticmethod def check_inputs(height, width, vae_scale_factor): @@ -142,13 +140,80 @@ class FluxProcessImagesInputStep(ModularPipelineBlocks): return components, state +class FluxKontextProcessImagesInputStep(ModularPipelineBlocks): + model_name = "flux-kontext" + + def __init__(self, _auto_resize=True): + self._auto_resize = _auto_resize + super().__init__() + + @property + def description(self) -> str: + return ( + "Image preprocess step for Flux Kontext. The preprocessed image goes to the VAE.\n" + "Kontext works as a T2I model, too, in case no input image is provided." + ) + + @property + def expected_components(self) -> List[ComponentSpec]: + return [ + ComponentSpec( + "image_processor", + VaeImageProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> List[InputParam]: + return [InputParam("image")] + + @property + def intermediate_outputs(self) -> List[OutputParam]: + return [OutputParam(name="processed_image")] + + @torch.no_grad() + def __call__(self, components: FluxModularPipeline, state: PipelineState): + from ...pipelines.flux.pipeline_flux_kontext import PREFERRED_KONTEXT_RESOLUTIONS + + block_state = self.get_block_state(state) + images = block_state.image + + if images is None: + block_state.processed_image = None + + else: + multiple_of = components.image_processor.config.vae_scale_factor + + if not is_valid_image_imagelist(images): + raise ValueError(f"Images must be image or list of images but are {type(images)}") + + if is_valid_image(images): + images = [images] + + img = images[0] + image_height, image_width = components.image_processor.get_default_height_width(img) + aspect_ratio = image_width / image_height + if self._auto_resize: + # Kontext is trained on specific resolutions, using one of them is recommended + _, image_width, image_height = min( + (abs(aspect_ratio - w / h), w, h) for w, h in PREFERRED_KONTEXT_RESOLUTIONS + ) + image_width = image_width // multiple_of * multiple_of + image_height = image_height // multiple_of * multiple_of + images = components.image_processor.resize(images, image_height, image_width) + block_state.processed_image = components.image_processor.preprocess(images, image_height, image_width) + + self.set_block_state(state, block_state) + return components, state + + class FluxVaeEncoderDynamicStep(ModularPipelineBlocks): model_name = "flux" def __init__( - self, - input_name: str = "processed_image", - output_name: str = "image_latents", + self, input_name: str = "processed_image", output_name: str = "image_latents", sample_mode: str = "sample" ): """Initialize a VAE encoder step for converting images to latent representations. @@ -160,6 +225,7 @@ class FluxVaeEncoderDynamicStep(ModularPipelineBlocks): Examples: "processed_image" or "processed_control_image" output_name (str, optional): Name of the output latent tensor. Defaults to "image_latents". Examples: "image_latents" or "control_image_latents" + sample_mode (str, optional): Sampling mode to be used. Examples: # Basic usage with default settings (includes image processor): # FluxImageVaeEncoderDynamicStep() @@ -170,6 +236,7 @@ class FluxVaeEncoderDynamicStep(ModularPipelineBlocks): """ self._image_input_name = input_name self._image_latents_output_name = output_name + self.sample_mode = sample_mode super().__init__() @property @@ -183,7 +250,7 @@ class FluxVaeEncoderDynamicStep(ModularPipelineBlocks): @property def inputs(self) -> List[InputParam]: - inputs = [InputParam(self._image_input_name, required=True), InputParam("generator")] + inputs = [InputParam(self._image_input_name), InputParam("generator")] return inputs @property @@ -199,16 +266,20 @@ class FluxVaeEncoderDynamicStep(ModularPipelineBlocks): @torch.no_grad() def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) - - device = components._execution_device - dtype = components.vae.dtype - image = getattr(block_state, self._image_input_name) - image = image.to(device=device, dtype=dtype) - # Encode image into latents - image_latents = encode_vae_image(image=image, vae=components.vae, generator=block_state.generator) - setattr(block_state, self._image_latents_output_name, image_latents) + if image is None: + setattr(block_state, self._image_latents_output_name, None) + else: + device = components._execution_device + dtype = components.vae.dtype + image = image.to(device=device, dtype=dtype) + + # Encode image into latents + image_latents = encode_vae_image( + image=image, vae=components.vae, generator=block_state.generator, sample_mode=self.sample_mode + ) + setattr(block_state, self._image_latents_output_name, image_latents) self.set_block_state(state, block_state) diff --git a/src/diffusers/modular_pipelines/flux/inputs.py b/src/diffusers/modular_pipelines/flux/inputs.py index f9192655d1..e1bc17f5ff 100644 --- a/src/diffusers/modular_pipelines/flux/inputs.py +++ b/src/diffusers/modular_pipelines/flux/inputs.py @@ -17,6 +17,7 @@ from typing import List import torch from ...pipelines import FluxPipeline +from ...utils import logging from ..modular_pipeline import ModularPipelineBlocks, PipelineState from ..modular_pipeline_utils import InputParam, OutputParam @@ -25,6 +26,9 @@ from ..qwenimage.inputs import calculate_dimension_from_latents, repeat_tensor_t from .modular_pipeline import FluxModularPipeline +logger = logging.get_logger(__name__) + + class FluxTextInputStep(ModularPipelineBlocks): model_name = "flux" @@ -234,3 +238,122 @@ class FluxInputsDynamicStep(ModularPipelineBlocks): self.set_block_state(state, block_state) return components, state + + +class FluxKontextInputsDynamicStep(FluxInputsDynamicStep): + model_name = "flux-kontext" + + def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + # Process image latent inputs (height/width calculation, patchify, and batch expansion) + for image_latent_input_name in self._image_latent_inputs: + image_latent_tensor = getattr(block_state, image_latent_input_name) + if image_latent_tensor is None: + continue + + # 1. Calculate height/width from latents + # Unlike the `FluxInputsDynamicStep`, we don't overwrite the `block.height` and `block.width` + height, width = calculate_dimension_from_latents(image_latent_tensor, components.vae_scale_factor) + if not hasattr(block_state, "image_height"): + block_state.image_height = height + if not hasattr(block_state, "image_width"): + block_state.image_width = width + + # 2. Patchify the image latent tensor + # TODO: Implement patchifier for Flux. + latent_height, latent_width = image_latent_tensor.shape[2:] + image_latent_tensor = FluxPipeline._pack_latents( + image_latent_tensor, block_state.batch_size, image_latent_tensor.shape[1], latent_height, latent_width + ) + + # 3. Expand batch size + image_latent_tensor = repeat_tensor_to_batch_size( + input_name=image_latent_input_name, + input_tensor=image_latent_tensor, + num_images_per_prompt=block_state.num_images_per_prompt, + batch_size=block_state.batch_size, + ) + + setattr(block_state, image_latent_input_name, image_latent_tensor) + + # Process additional batch inputs (only batch expansion) + for input_name in self._additional_batch_inputs: + input_tensor = getattr(block_state, input_name) + if input_tensor is None: + continue + + # Only expand batch size + input_tensor = repeat_tensor_to_batch_size( + input_name=input_name, + input_tensor=input_tensor, + num_images_per_prompt=block_state.num_images_per_prompt, + batch_size=block_state.batch_size, + ) + + setattr(block_state, input_name, input_tensor) + + self.set_block_state(state, block_state) + return components, state + + +class FluxKontextSetResolutionStep(ModularPipelineBlocks): + model_name = "flux-kontext" + + def description(self): + return ( + "Determines the height and width to be used during the subsequent computations.\n" + "It should always be placed _before_ the latent preparation step." + ) + + @property + def inputs(self) -> List[InputParam]: + inputs = [ + InputParam(name="height"), + InputParam(name="width"), + InputParam(name="max_area", type_hint=int, default=1024**2), + ] + return inputs + + @property + def intermediate_outputs(self) -> List[OutputParam]: + return [ + OutputParam(name="height", type_hint=int, description="The height of the initial noisy latents"), + OutputParam(name="width", type_hint=int, description="The width of the initial noisy latents"), + ] + + @staticmethod + def check_inputs(height, width, vae_scale_factor): + if height is not None and height % (vae_scale_factor * 2) != 0: + raise ValueError(f"Height must be divisible by {vae_scale_factor * 2} but is {height}") + + if width is not None and width % (vae_scale_factor * 2) != 0: + raise ValueError(f"Width must be divisible by {vae_scale_factor * 2} but is {width}") + + def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + height = block_state.height or components.default_height + width = block_state.width or components.default_width + self.check_inputs(height, width, components.vae_scale_factor) + + original_height, original_width = height, width + max_area = block_state.max_area + aspect_ratio = width / height + width = round((max_area * aspect_ratio) ** 0.5) + height = round((max_area / aspect_ratio) ** 0.5) + + multiple_of = components.vae_scale_factor * 2 + width = width // multiple_of * multiple_of + height = height // multiple_of * multiple_of + + if height != original_height or width != original_width: + logger.warning( + f"Generation `height` and `width` have been adjusted to {height} and {width} to fit the model requirements." + ) + + block_state.height = height + block_state.width = width + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/flux/modular_blocks.py b/src/diffusers/modular_pipelines/flux/modular_blocks.py index b40dfe1762..a80bc2a5f7 100644 --- a/src/diffusers/modular_pipelines/flux/modular_blocks.py +++ b/src/diffusers/modular_pipelines/flux/modular_blocks.py @@ -18,14 +18,25 @@ from ..modular_pipeline_utils import InsertableDict from .before_denoise import ( FluxImg2ImgPrepareLatentsStep, FluxImg2ImgSetTimestepsStep, + FluxKontextRoPEInputsStep, FluxPrepareLatentsStep, FluxRoPEInputsStep, FluxSetTimestepsStep, ) from .decoders import FluxDecodeStep -from .denoise import FluxDenoiseStep -from .encoders import FluxProcessImagesInputStep, FluxTextEncoderStep, FluxVaeEncoderDynamicStep -from .inputs import FluxInputsDynamicStep, FluxTextInputStep +from .denoise import FluxDenoiseStep, FluxKontextDenoiseStep +from .encoders import ( + FluxKontextProcessImagesInputStep, + FluxProcessImagesInputStep, + FluxTextEncoderStep, + FluxVaeEncoderDynamicStep, +) +from .inputs import ( + FluxInputsDynamicStep, + FluxKontextInputsDynamicStep, + FluxKontextSetResolutionStep, + FluxTextInputStep, +) logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -33,10 +44,7 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name # vae encoder (run before before_denoise) FluxImg2ImgVaeEncoderBlocks = InsertableDict( - [ - ("preprocess", FluxProcessImagesInputStep()), - ("encode", FluxVaeEncoderDynamicStep()), - ] + [("preprocess", FluxProcessImagesInputStep()), ("encode", FluxVaeEncoderDynamicStep())] ) @@ -66,6 +74,39 @@ class FluxAutoVaeEncoderStep(AutoPipelineBlocks): ) +# Flux Kontext vae encoder (run before before_denoise) + +FluxKontextVaeEncoderBlocks = InsertableDict( + [("preprocess", FluxKontextProcessImagesInputStep()), ("encode", FluxVaeEncoderDynamicStep(sample_mode="argmax"))] +) + + +class FluxKontextVaeEncoderStep(SequentialPipelineBlocks): + model_name = "flux-kontext" + + block_classes = FluxKontextVaeEncoderBlocks.values() + block_names = FluxKontextVaeEncoderBlocks.keys() + + @property + def description(self) -> str: + return "Vae encoder step that preprocess andencode the image inputs into their latent representations." + + +class FluxKontextAutoVaeEncoderStep(AutoPipelineBlocks): + block_classes = [FluxKontextVaeEncoderStep] + block_names = ["img2img"] + block_trigger_inputs = ["image"] + + @property + def description(self): + return ( + "Vae encoder step that encode the image inputs into their latent representations.\n" + + "This is an auto pipeline block that works for img2img tasks.\n" + + " - `FluxKontextVaeEncoderStep` (img2img) is used when only `image` is provided." + + " - if `image` is not provided, step will be skipped." + ) + + # before_denoise: text2img FluxBeforeDenoiseBlocks = InsertableDict( [ @@ -107,6 +148,7 @@ class FluxImg2ImgBeforeDenoiseStep(SequentialPipelineBlocks): # before_denoise: all task (text2img, img2img) class FluxAutoBeforeDenoiseStep(AutoPipelineBlocks): + model_name = "flux-kontext" block_classes = [FluxImg2ImgBeforeDenoiseStep, FluxBeforeDenoiseStep] block_names = ["img2img", "text2image"] block_trigger_inputs = ["image_latents", None] @@ -121,6 +163,44 @@ class FluxAutoBeforeDenoiseStep(AutoPipelineBlocks): ) +# before_denoise: FluxKontext + +FluxKontextBeforeDenoiseBlocks = InsertableDict( + [ + ("prepare_latents", FluxPrepareLatentsStep()), + ("set_timesteps", FluxSetTimestepsStep()), + ("prepare_rope_inputs", FluxKontextRoPEInputsStep()), + ] +) + + +class FluxKontextBeforeDenoiseStep(SequentialPipelineBlocks): + block_classes = FluxKontextBeforeDenoiseBlocks.values() + block_names = FluxKontextBeforeDenoiseBlocks.keys() + + @property + def description(self): + return ( + "Before denoise step that prepare the inputs for the denoise step\n" + "for img2img/text2img task for Flux Kontext." + ) + + +class FluxKontextAutoBeforeDenoiseStep(AutoPipelineBlocks): + block_classes = [FluxKontextBeforeDenoiseStep, FluxBeforeDenoiseStep] + block_names = ["img2img", "text2image"] + block_trigger_inputs = ["image_latents", None] + + @property + def description(self): + return ( + "Before denoise step that prepare the inputs for the denoise step.\n" + + "This is an auto pipeline block that works for text2image.\n" + + " - `FluxBeforeDenoiseStep` (text2image) is used.\n" + + " - `FluxKontextBeforeDenoiseStep` (img2img) is used when only `image_latents` is provided.\n" + ) + + # denoise: text2image class FluxAutoDenoiseStep(AutoPipelineBlocks): block_classes = [FluxDenoiseStep] @@ -136,6 +216,23 @@ class FluxAutoDenoiseStep(AutoPipelineBlocks): ) +# denoise: Flux Kontext + + +class FluxKontextAutoDenoiseStep(AutoPipelineBlocks): + block_classes = [FluxKontextDenoiseStep] + block_names = ["denoise"] + block_trigger_inputs = [None] + + @property + def description(self) -> str: + return ( + "Denoise step that iteratively denoise the latents for Flux Kontext. " + "This is a auto pipeline block that works for text2image and img2img tasks." + " - `FluxDenoiseStep` (denoise) for text2image and img2img tasks." + ) + + # decode: all task (text2img, img2img) class FluxAutoDecodeStep(AutoPipelineBlocks): block_classes = [FluxDecodeStep] @@ -165,7 +262,7 @@ class FluxImg2ImgInputStep(SequentialPipelineBlocks): " - update height/width based `image_latents`, patchify `image_latents`." -class FluxImageAutoInputStep(AutoPipelineBlocks): +class FluxAutoInputStep(AutoPipelineBlocks): block_classes = [FluxImg2ImgInputStep, FluxTextInputStep] block_names = ["img2img", "text2image"] block_trigger_inputs = ["image_latents", None] @@ -180,16 +277,59 @@ class FluxImageAutoInputStep(AutoPipelineBlocks): ) +# inputs: Flux Kontext + +FluxKontextBlocks = InsertableDict( + [ + ("set_resolution", FluxKontextSetResolutionStep()), + ("text_inputs", FluxTextInputStep()), + ("additional_inputs", FluxKontextInputsDynamicStep()), + ] +) + + +class FluxKontextInputStep(SequentialPipelineBlocks): + model_name = "flux-kontext" + block_classes = FluxKontextBlocks.values() + block_names = FluxKontextBlocks.keys() + + @property + def description(self): + return ( + "Input step that prepares the inputs for the both text2img and img2img denoising step. It:\n" + " - make sure the text embeddings have consistent batch size as well as the additional inputs (`image_latents`).\n" + " - update height/width based `image_latents`, patchify `image_latents`." + ) + + +class FluxKontextAutoInputStep(AutoPipelineBlocks): + block_classes = [FluxKontextInputStep, FluxTextInputStep] + # block_classes = [FluxKontextInputStep] + block_names = ["img2img", "text2img"] + # block_names = ["img2img"] + block_trigger_inputs = ["image_latents", None] + # block_trigger_inputs = ["image_latents"] + + @property + def description(self): + return ( + "Input step that standardize the inputs for the denoising step, e.g. make sure inputs have consistent batch size, and patchified. \n" + " This is an auto pipeline block that works for text2image/img2img tasks.\n" + + " - `FluxKontextInputStep` (img2img) is used when `image_latents` is provided.\n" + + " - `FluxKontextInputStep` is also capable of handling text2image task when `image_latent` isn't present." + ) + + class FluxCoreDenoiseStep(SequentialPipelineBlocks): model_name = "flux" - block_classes = [FluxImageAutoInputStep, FluxAutoBeforeDenoiseStep, FluxAutoDenoiseStep] + block_classes = [FluxAutoInputStep, FluxAutoBeforeDenoiseStep, FluxAutoDenoiseStep] block_names = ["input", "before_denoise", "denoise"] @property def description(self): return ( "Core step that performs the denoising process. \n" - + " - `FluxImageAutoInputStep` (input) standardizes the inputs for the denoising step.\n" + + " - `FluxAutoInputStep` (input) standardizes the inputs for the denoising step.\n" + " - `FluxAutoBeforeDenoiseStep` (before_denoise) prepares the inputs for the denoising step.\n" + " - `FluxAutoDenoiseStep` (denoise) iteratively denoises the latents.\n" + "This step supports text-to-image and image-to-image tasks for Flux:\n" @@ -198,6 +338,24 @@ class FluxCoreDenoiseStep(SequentialPipelineBlocks): ) +class FluxKontextCoreDenoiseStep(SequentialPipelineBlocks): + model_name = "flux-kontext" + block_classes = [FluxKontextAutoInputStep, FluxKontextAutoBeforeDenoiseStep, FluxKontextAutoDenoiseStep] + block_names = ["input", "before_denoise", "denoise"] + + @property + def description(self): + return ( + "Core step that performs the denoising process. \n" + + " - `FluxKontextAutoInputStep` (input) standardizes the inputs for the denoising step.\n" + + " - `FluxKontextAutoBeforeDenoiseStep` (before_denoise) prepares the inputs for the denoising step.\n" + + " - `FluxKontextAutoDenoiseStep` (denoise) iteratively denoises the latents.\n" + + "This step supports text-to-image and image-to-image tasks for Flux:\n" + + " - for image-to-image generation, you need to provide `image_latents`\n" + + " - for text-to-image generation, all you need to provide is prompt embeddings." + ) + + # Auto blocks (text2image and img2img) AUTO_BLOCKS = InsertableDict( [ @@ -208,6 +366,15 @@ AUTO_BLOCKS = InsertableDict( ] ) +AUTO_BLOCKS_KONTEXT = InsertableDict( + [ + ("text_encoder", FluxTextEncoderStep()), + ("image_encoder", FluxKontextAutoVaeEncoderStep()), + ("denoise", FluxKontextCoreDenoiseStep()), + ("decode", FluxDecodeStep()), + ] +) + class FluxAutoBlocks(SequentialPipelineBlocks): model_name = "flux" @@ -224,6 +391,13 @@ class FluxAutoBlocks(SequentialPipelineBlocks): ) +class FluxKontextAutoBlocks(FluxAutoBlocks): + model_name = "flux-kontext" + + block_classes = AUTO_BLOCKS_KONTEXT.values() + block_names = AUTO_BLOCKS_KONTEXT.keys() + + TEXT2IMAGE_BLOCKS = InsertableDict( [ ("text_encoder", FluxTextEncoderStep()), @@ -250,4 +424,23 @@ IMAGE2IMAGE_BLOCKS = InsertableDict( ] ) -ALL_BLOCKS = {"text2image": TEXT2IMAGE_BLOCKS, "img2img": IMAGE2IMAGE_BLOCKS, "auto": AUTO_BLOCKS} +FLUX_KONTEXT_BLOCKS = InsertableDict( + [ + ("text_encoder", FluxTextEncoderStep()), + ("vae_encoder", FluxVaeEncoderDynamicStep(sample_mode="argmax")), + ("input", FluxKontextInputStep()), + ("prepare_latents", FluxPrepareLatentsStep()), + ("set_timesteps", FluxSetTimestepsStep()), + ("prepare_rope_inputs", FluxKontextRoPEInputsStep()), + ("denoise", FluxKontextDenoiseStep()), + ("decode", FluxDecodeStep()), + ] +) + +ALL_BLOCKS = { + "text2image": TEXT2IMAGE_BLOCKS, + "img2img": IMAGE2IMAGE_BLOCKS, + "auto": AUTO_BLOCKS, + "auto_kontext": AUTO_BLOCKS_KONTEXT, + "kontext": FLUX_KONTEXT_BLOCKS, +} diff --git a/src/diffusers/modular_pipelines/flux/modular_pipeline.py b/src/diffusers/modular_pipelines/flux/modular_pipeline.py index 563b033343..d8158f5d4f 100644 --- a/src/diffusers/modular_pipelines/flux/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/flux/modular_pipeline.py @@ -55,3 +55,13 @@ class FluxModularPipeline(ModularPipeline, FluxLoraLoaderMixin, TextualInversion if getattr(self, "transformer", None): num_channels_latents = self.transformer.config.in_channels // 4 return num_channels_latents + + +class FluxKontextModularPipeline(FluxModularPipeline): + """ + A ModularPipeline for Flux Kontext. + + > [!WARNING] > This is an experimental feature and is likely to change in the future. + """ + + default_blocks_name = "FluxKontextAutoBlocks" diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index e543bf0bb3..cfbca48a98 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -57,6 +57,7 @@ MODULAR_PIPELINE_MAPPING = OrderedDict( ("stable-diffusion-xl", "StableDiffusionXLModularPipeline"), ("wan", "WanModularPipeline"), ("flux", "FluxModularPipeline"), + ("flux-kontext", "FluxKontextModularPipeline"), ("qwenimage", "QwenImageModularPipeline"), ("qwenimage-edit", "QwenImageEditModularPipeline"), ("qwenimage-edit-plus", "QwenImageEditPlusModularPipeline"), diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index cf80377964..9ed6250452 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -17,6 +17,36 @@ class FluxAutoBlocks(metaclass=DummyObject): requires_backends(cls, ["torch", "transformers"]) +class FluxKontextAutoBlocks(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class FluxKontextModularPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class FluxModularPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] From 8abc7aeb715c0149ee0a9982b2d608ce97f55215 Mon Sep 17 00:00:00 2001 From: Steven Liu <59462357+stevhliu@users.noreply.github.com> Date: Fri, 10 Oct 2025 19:43:30 -0700 Subject: [PATCH 061/123] [docs] Fix syntax (#12464) * fix syntax * fix * style * fix --- docs/source/en/api/pipelines/marigold.md | 2 +- .../marigold/pipeline_marigold_depth.py | 11 +++++------ .../marigold/pipeline_marigold_intrinsics.py | 16 ++++++++-------- .../marigold/pipeline_marigold_normals.py | 11 +++++------ 4 files changed, 19 insertions(+), 21 deletions(-) diff --git a/docs/source/en/api/pipelines/marigold.md b/docs/source/en/api/pipelines/marigold.md index 81e103afeb..bb6e94de33 100644 --- a/docs/source/en/api/pipelines/marigold.md +++ b/docs/source/en/api/pipelines/marigold.md @@ -75,7 +75,7 @@ The following is a summary of the recommended checkpoints, all of which produce | [prs-eth/marigold-depth-v1-1](https://huggingface.co/prs-eth/marigold-depth-v1-1) | Depth | Affine-invariant depth prediction assigns each pixel a value between 0 (near plane) and 1 (far plane), with both planes determined by the model during inference. | | [prs-eth/marigold-normals-v0-1](https://huggingface.co/prs-eth/marigold-normals-v0-1) | Normals | The surface normals predictions are unit-length 3D vectors in the screen space camera, with values in the range from -1 to 1. | | [prs-eth/marigold-iid-appearance-v1-1](https://huggingface.co/prs-eth/marigold-iid-appearance-v1-1) | Intrinsics | InteriorVerse decomposition is comprised of Albedo and two BRDF material properties: Roughness and Metallicity. | -| [prs-eth/marigold-iid-lighting-v1-1](https://huggingface.co/prs-eth/marigold-iid-lighting-v1-1) | Intrinsics | HyperSim decomposition of an image  \\(I\\)  is comprised of Albedo  \\(A\\), Diffuse shading  \\(S\\), and Non-diffuse residual  \\(R\\):  \\(I = A*S+R\\). | +| [prs-eth/marigold-iid-lighting-v1-1](https://huggingface.co/prs-eth/marigold-iid-lighting-v1-1) | Intrinsics | HyperSim decomposition of an image $I$ is comprised of Albedo $A$, Diffuse shading $S$, and Non-diffuse residual $R$: $I = A*S+R$. | > [!TIP] > Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff diff --git a/src/diffusers/pipelines/marigold/pipeline_marigold_depth.py b/src/diffusers/pipelines/marigold/pipeline_marigold_depth.py index da991aefbd..92ec16fd45 100644 --- a/src/diffusers/pipelines/marigold/pipeline_marigold_depth.py +++ b/src/diffusers/pipelines/marigold/pipeline_marigold_depth.py @@ -86,15 +86,14 @@ class MarigoldDepthOutput(BaseOutput): Args: prediction (`np.ndarray`, `torch.Tensor`): - Predicted depth maps with values in the range [0, 1]. The shape is $numimages \times 1 \times height \times - width$ for `torch.Tensor` or $numimages \times height \times width \times 1$ for `np.ndarray`. + Predicted depth maps with values in the range [0, 1]. The shape is `numimages × 1 × height × width` for + `torch.Tensor` or `numimages × height × width × 1` for `np.ndarray`. uncertainty (`None`, `np.ndarray`, `torch.Tensor`): - Uncertainty maps computed from the ensemble, with values in the range [0, 1]. The shape is $numimages - \times 1 \times height \times width$ for `torch.Tensor` or $numimages \times height \times width \times 1$ - for `np.ndarray`. + Uncertainty maps computed from the ensemble, with values in the range [0, 1]. The shape is `numimages × 1 × + height × width` for `torch.Tensor` or `numimages × height × width × 1` for `np.ndarray`. latent (`None`, `torch.Tensor`): Latent features corresponding to the predictions, compatible with the `latents` argument of the pipeline. - The shape is $numimages * numensemble \times 4 \times latentheight \times latentwidth$. + The shape is `numimages * numensemble × 4 × latentheight × latentwidth`. """ prediction: Union[np.ndarray, torch.Tensor] diff --git a/src/diffusers/pipelines/marigold/pipeline_marigold_intrinsics.py b/src/diffusers/pipelines/marigold/pipeline_marigold_intrinsics.py index c809de18f4..bef9ca77c7 100644 --- a/src/diffusers/pipelines/marigold/pipeline_marigold_intrinsics.py +++ b/src/diffusers/pipelines/marigold/pipeline_marigold_intrinsics.py @@ -99,17 +99,17 @@ class MarigoldIntrinsicsOutput(BaseOutput): Args: prediction (`np.ndarray`, `torch.Tensor`): - Predicted image intrinsics with values in the range [0, 1]. The shape is $(numimages * numtargets) \times 3 - \times height \times width$ for `torch.Tensor` or $(numimages * numtargets) \times height \times width - \times 3$ for `np.ndarray`, where `numtargets` corresponds to the number of predicted target modalities of - the intrinsic image decomposition. + Predicted image intrinsics with values in the range [0, 1]. The shape is `(numimages * numtargets) × 3 × + height × width` for `torch.Tensor` or `(numimages * numtargets) × height × width × 3` for `np.ndarray`, + where `numtargets` corresponds to the number of predicted target modalities of the intrinsic image + decomposition. uncertainty (`None`, `np.ndarray`, `torch.Tensor`): - Uncertainty maps computed from the ensemble, with values in the range [0, 1]. The shape is $(numimages * - numtargets) \times 3 \times height \times width$ for `torch.Tensor` or $(numimages * numtargets) \times - height \times width \times 3$ for `np.ndarray`. + Uncertainty maps computed from the ensemble, with values in the range [0, 1]. The shape is `(numimages * + numtargets) × 3 × height × width` for `torch.Tensor` or `(numimages * numtargets) × height × width × 3` for + `np.ndarray`. latent (`None`, `torch.Tensor`): Latent features corresponding to the predictions, compatible with the `latents` argument of the pipeline. - The shape is $(numimages * numensemble) \times (numtargets * 4) \times latentheight \times latentwidth$. + The shape is `(numimages * numensemble) × (numtargets * 4) × latentheight × latentwidth`. """ prediction: Union[np.ndarray, torch.Tensor] diff --git a/src/diffusers/pipelines/marigold/pipeline_marigold_normals.py b/src/diffusers/pipelines/marigold/pipeline_marigold_normals.py index 192ed590a4..485a39c995 100644 --- a/src/diffusers/pipelines/marigold/pipeline_marigold_normals.py +++ b/src/diffusers/pipelines/marigold/pipeline_marigold_normals.py @@ -81,15 +81,14 @@ class MarigoldNormalsOutput(BaseOutput): Args: prediction (`np.ndarray`, `torch.Tensor`): - Predicted normals with values in the range [-1, 1]. The shape is $numimages \times 3 \times height \times - width$ for `torch.Tensor` or $numimages \times height \times width \times 3$ for `np.ndarray`. + Predicted normals with values in the range [-1, 1]. The shape is `numimages × 3 × height × width` for + `torch.Tensor` or `numimages × height × width × 3` for `np.ndarray`. uncertainty (`None`, `np.ndarray`, `torch.Tensor`): - Uncertainty maps computed from the ensemble, with values in the range [0, 1]. The shape is $numimages - \times 1 \times height \times width$ for `torch.Tensor` or $numimages \times height \times width \times 1$ - for `np.ndarray`. + Uncertainty maps computed from the ensemble, with values in the range [0, 1]. The shape is `numimages × 1 × + height × width` for `torch.Tensor` or `numimages × height × width × 1` for `np.ndarray`. latent (`None`, `torch.Tensor`): Latent features corresponding to the predictions, compatible with the `latents` argument of the pipeline. - The shape is $numimages * numensemble \times 4 \times latentheight \times latentwidth$. + The shape is `numimages * numensemble × 4 × latentheight × latentwidth`. """ prediction: Union[np.ndarray, torch.Tensor] From fa468c5d576a42f50ef467a4dc24d0326b25f656 Mon Sep 17 00:00:00 2001 From: Manith Ratnayake <144333591+Manith-Ratnayake@users.noreply.github.com> Date: Mon, 13 Oct 2025 21:27:46 +0530 Subject: [PATCH 062/123] docs: api-pipelines-qwenimage typo fix (#12461) --- docs/source/en/api/pipelines/qwenimage.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/en/api/pipelines/qwenimage.md b/docs/source/en/api/pipelines/qwenimage.md index 27cc4802f6..b3dd3dd936 100644 --- a/docs/source/en/api/pipelines/qwenimage.md +++ b/docs/source/en/api/pipelines/qwenimage.md @@ -109,7 +109,7 @@ image_1 = load_image("https://huggingface.co/datasets/huggingface/documentation- image_2 = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/peng.png") image = pipe( image=[image_1, image_2], - prompt="put the penguin and the cat at a game show called "Qwen Edit Plus Games"", + prompt='''put the penguin and the cat at a game show called "Qwen Edit Plus Games"''', num_inference_steps=50 ).images[0] ``` From a4bc8454783dbae3be9cf840f074b961a558aba5 Mon Sep 17 00:00:00 2001 From: Meatfucker <74834323+Meatfucker@users.noreply.github.com> Date: Tue, 14 Oct 2025 13:43:21 -0400 Subject: [PATCH 063/123] Fix missing load_video documentation and load_video import in WanVideoToVideoPipeline example code (#12472) * Update utilities.md Update missing load_video documentation * Update pipeline_wan_video2video.py Fix missing load_video import in example code --- docs/source/en/api/utilities.md | 4 ++++ src/diffusers/pipelines/wan/pipeline_wan_video2video.py | 2 +- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/docs/source/en/api/utilities.md b/docs/source/en/api/utilities.md index 576c331834..45a3c5372a 100644 --- a/docs/source/en/api/utilities.md +++ b/docs/source/en/api/utilities.md @@ -26,6 +26,10 @@ Utility and helper functions for working with 🤗 Diffusers. [[autodoc]] utils.load_image +## load_video + +[[autodoc]] utils.load_video + ## export_to_gif [[autodoc]] utils.export_to_gif diff --git a/src/diffusers/pipelines/wan/pipeline_wan_video2video.py b/src/diffusers/pipelines/wan/pipeline_wan_video2video.py index 1a2d2e9c22..a976126da7 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan_video2video.py +++ b/src/diffusers/pipelines/wan/pipeline_wan_video2video.py @@ -49,7 +49,7 @@ EXAMPLE_DOC_STRING = """ Examples: ```python >>> import torch - >>> from diffusers.utils import export_to_video + >>> from diffusers.utils import export_to_video, load_video >>> from diffusers import AutoencoderKLWan, WanVideoToVideoPipeline >>> from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler From 3eb40786cad189f191f5dacca1e84a11ed015dd6 Mon Sep 17 00:00:00 2001 From: Steven Liu <59462357+stevhliu@users.noreply.github.com> Date: Tue, 14 Oct 2025 13:53:56 -0700 Subject: [PATCH 064/123] [docs] Prompting (#12312) * init * fix * batch inf * feedback * update --- .gitignore | 3 + docs/source/en/_toctree.yml | 2 +- .../en/using-diffusers/batched_inference.md | 175 ++----- .../en/using-diffusers/weighted_prompts.md | 435 +++--------------- 4 files changed, 100 insertions(+), 515 deletions(-) diff --git a/.gitignore b/.gitignore index 15617d5fdc..a55026febd 100644 --- a/.gitignore +++ b/.gitignore @@ -125,6 +125,9 @@ dmypy.json .vs .vscode +# Cursor +.cursor + # Pycharm .idea diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index fb4fdf2098..848e38079e 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -49,7 +49,7 @@ isExpanded: false sections: - local: using-diffusers/weighted_prompts - title: Prompt techniques + title: Prompting - local: using-diffusers/create_a_server title: Create a server - local: using-diffusers/batched_inference diff --git a/docs/source/en/using-diffusers/batched_inference.md b/docs/source/en/using-diffusers/batched_inference.md index b5e55c27ca..cdb16ac121 100644 --- a/docs/source/en/using-diffusers/batched_inference.md +++ b/docs/source/en/using-diffusers/batched_inference.md @@ -16,24 +16,24 @@ Batch inference processes multiple prompts at a time to increase throughput. It The downside is increased latency because you must wait for the entire batch to complete, and more GPU memory is required for large batches. - - - -For text-to-image, pass a list of prompts to the pipeline. +For text-to-image, pass a list of prompts to the pipeline and for image-to-image, pass a list of images and prompts to the pipeline. The example below demonstrates batched text-to-image inference. ```py import torch +import matplotlib.pyplot as plt from diffusers import DiffusionPipeline pipeline = DiffusionPipeline.from_pretrained( "stabilityai/stable-diffusion-xl-base-1.0", - torch_dtype=torch.float16 -).to("cuda") + torch_dtype=torch.float16, + device_map="cuda" +) prompts = [ - "cinematic photo of A beautiful sunset over mountains, 35mm photograph, film, professional, 4k, highly detailed", - "cinematic film still of a cat basking in the sun on a roof in Turkey, highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain", - "pixel-art a cozy coffee shop interior, low-res, blocky, pixel art style, 8-bit graphics" + "Cinematic shot of a cozy coffee shop interior, warm pastel light streaming through a window where a cat rests. Shallow depth of field, glowing cups in soft focus, dreamy lofi-inspired mood, nostalgic tones, framed like a quiet film scene.", + "Polaroid-style photograph of a cozy coffee shop interior, bathed in warm pastel light. A cat sits on the windowsill near steaming mugs. Soft, slightly faded tones and dreamy blur evoke nostalgia, a lofi mood, and the intimate, imperfect charm of instant film.", + "Soft watercolor illustration of a cozy coffee shop interior, pastel washes of color filling the space. A cat rests peacefully on the windowsill as warm light glows through. Gentle brushstrokes create a dreamy, lofi-inspired atmosphere with whimsical textures and nostalgic calm.", + "Isometric pixel-art illustration of a cozy coffee shop interior in detailed 8-bit style. Warm pastel light fills the space as a cat rests on the windowsill. Blocky furniture and tiny mugs add charm, low-res retro graphics enhance the nostalgic, lofi-inspired game aesthetic." ] images = pipeline( @@ -52,6 +52,10 @@ plt.tight_layout() plt.show() ``` +
+ +
+ To generate multiple variations of one prompt, use the `num_images_per_prompt` argument. ```py @@ -61,11 +65,18 @@ from diffusers import DiffusionPipeline pipeline = DiffusionPipeline.from_pretrained( "stabilityai/stable-diffusion-xl-base-1.0", - torch_dtype=torch.float16 -).to("cuda") + torch_dtype=torch.float16, + device_map="cuda" +) + +prompt=""" +Isometric pixel-art illustration of a cozy coffee shop interior in detailed 8-bit style. Warm pastel light fills the +space as a cat rests on the windowsill. Blocky furniture and tiny mugs add charm, low-res retro graphics enhance the +nostalgic, lofi-inspired game aesthetic. +""" images = pipeline( - prompt="pixel-art a cozy coffee shop interior, low-res, blocky, pixel art style, 8-bit graphics", + prompt=prompt, num_images_per_prompt=4 ).images @@ -81,6 +92,10 @@ plt.tight_layout() plt.show() ``` +
+ +
+ Combine both approaches to generate different variations of different prompts. ```py @@ -89,7 +104,7 @@ images = pipeline( num_images_per_prompt=2, ).images -fig, axes = plt.subplots(2, 2, figsize=(12, 12)) +fig, axes = plt.subplots(2, 4, figsize=(12, 12)) axes = axes.flatten() for i, image in enumerate(images): @@ -101,126 +116,18 @@ plt.tight_layout() plt.show() ``` - - - -For image-to-image, pass a list of input images and prompts to the pipeline. - -```py -import torch -from diffusers.utils import load_image -from diffusers import DiffusionPipeline - -pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", - torch_dtype=torch.float16 -).to("cuda") - -input_images = [ - load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png"), - load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png"), - load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/detail-prompt.png") -] - -prompts = [ - "cinematic photo of a beautiful sunset over mountains, 35mm photograph, film, professional, 4k, highly detailed", - "cinematic film still of a cat basking in the sun on a roof in Turkey, highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain", - "pixel-art a cozy coffee shop interior, low-res, blocky, pixel art style, 8-bit graphics" -] - -images = pipeline( - prompt=prompts, - image=input_images, - guidance_scale=8.0, - strength=0.5 -).images - -fig, axes = plt.subplots(2, 2, figsize=(12, 12)) -axes = axes.flatten() - -for i, image in enumerate(images): - axes[i].imshow(image) - axes[i].set_title(f"Image {i+1}") - axes[i].axis('off') - -plt.tight_layout() -plt.show() -``` - -To generate multiple variations of one prompt, use the `num_images_per_prompt` argument. - -```py -import torch -import matplotlib.pyplot as plt -from diffusers.utils import load_image -from diffusers import DiffusionPipeline - -pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", - torch_dtype=torch.float16 -).to("cuda") - -input_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/detail-prompt.png") - -images = pipeline( - prompt="pixel-art a cozy coffee shop interior, low-res, blocky, pixel art style, 8-bit graphics", - image=input_image, - num_images_per_prompt=4 -).images - -fig, axes = plt.subplots(2, 2, figsize=(12, 12)) -axes = axes.flatten() - -for i, image in enumerate(images): - axes[i].imshow(image) - axes[i].set_title(f"Image {i+1}") - axes[i].axis('off') - -plt.tight_layout() -plt.show() -``` - -Combine both approaches to generate different variations of different prompts. - -```py -input_images = [ - load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png"), - load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/detail-prompt.png") -] - -prompts = [ - "cinematic film still of a cat basking in the sun on a roof in Turkey, highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain", - "pixel-art a cozy coffee shop interior, low-res, blocky, pixel art style, 8-bit graphics" -] - -images = pipeline( - prompt=prompts, - image=input_images, - num_images_per_prompt=2, -).images - -fig, axes = plt.subplots(2, 2, figsize=(12, 12)) -axes = axes.flatten() - -for i, image in enumerate(images): - axes[i].imshow(image) - axes[i].set_title(f"Image {i+1}") - axes[i].axis('off') - -plt.tight_layout() -plt.show() -``` - - - +
+ +
## Deterministic generation Enable reproducible batch generation by passing a list of [Generator’s](https://pytorch.org/docs/stable/generated/torch.Generator.html) to the pipeline and tie each `Generator` to a seed to reuse it. -Use a list comprehension to iterate over the batch size specified in `range()` to create a unique `Generator` object for each image in the batch. +> [!TIP] +> Refer to the [Reproducibility](./reusing_seeds) docs to learn more about deterministic algorithms and the `Generator` object. -Don't multiply the `Generator` by the batch size because that only creates one `Generator` object that is used sequentially for each image in the batch. +Use a list comprehension to iterate over the batch size specified in `range()` to create a unique `Generator` object for each image in the batch. Don't multiply the `Generator` by the batch size because that only creates one `Generator` object that is used sequentially for each image in the batch. ```py generator = [torch.Generator(device="cuda").manual_seed(0)] * 3 @@ -234,14 +141,16 @@ from diffusers import DiffusionPipeline pipeline = DiffusionPipeline.from_pretrained( "stabilityai/stable-diffusion-xl-base-1.0", - torch_dtype=torch.float16 -).to("cuda") + torch_dtype=torch.float16, + device_map="cuda" +) generator = [torch.Generator(device="cuda").manual_seed(i) for i in range(3)] prompts = [ - "cinematic photo of A beautiful sunset over mountains, 35mm photograph, film, professional, 4k, highly detailed", - "cinematic film still of a cat basking in the sun on a roof in Turkey, highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain", - "pixel-art a cozy coffee shop interior, low-res, blocky, pixel art style, 8-bit graphics" + "Cinematic shot of a cozy coffee shop interior, warm pastel light streaming through a window where a cat rests. Shallow depth of field, glowing cups in soft focus, dreamy lofi-inspired mood, nostalgic tones, framed like a quiet film scene.", + "Polaroid-style photograph of a cozy coffee shop interior, bathed in warm pastel light. A cat sits on the windowsill near steaming mugs. Soft, slightly faded tones and dreamy blur evoke nostalgia, a lofi mood, and the intimate, imperfect charm of instant film.", + "Soft watercolor illustration of a cozy coffee shop interior, pastel washes of color filling the space. A cat rests peacefully on the windowsill as warm light glows through. Gentle brushstrokes create a dreamy, lofi-inspired atmosphere with whimsical textures and nostalgic calm.", + "Isometric pixel-art illustration of a cozy coffee shop interior in detailed 8-bit style. Warm pastel light fills the space as a cat rests on the windowsill. Blocky furniture and tiny mugs add charm, low-res retro graphics enhance the nostalgic, lofi-inspired game aesthetic." ] images = pipeline( @@ -261,4 +170,4 @@ plt.tight_layout() plt.show() ``` -You can use this to iteratively select an image associated with a seed and then improve on it by crafting a more detailed prompt. \ No newline at end of file +You can use this to select an image associated with a seed and iteratively improve on it by crafting a more detailed prompt. \ No newline at end of file diff --git a/docs/source/en/using-diffusers/weighted_prompts.md b/docs/source/en/using-diffusers/weighted_prompts.md index b45568ac4d..f89ebfe4a2 100644 --- a/docs/source/en/using-diffusers/weighted_prompts.md +++ b/docs/source/en/using-diffusers/weighted_prompts.md @@ -10,423 +10,96 @@ an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express o specific language governing permissions and limitations under the License. --> -# Prompt techniques - [[open-in-colab]] -Prompts are important because they describe what you want a diffusion model to generate. The best prompts are detailed, specific, and well-structured to help the model realize your vision. But crafting a great prompt takes time and effort and sometimes it may not be enough because language and words can be imprecise. This is where you need to boost your prompt with other techniques, such as prompt enhancing and prompt weighting, to get the results you want. +# Prompting -This guide will show you how you can use these prompt techniques to generate high-quality images with lower effort and adjust the weight of certain keywords in a prompt. +Prompts describes what a model should generate. Good prompts are detailed, specific, and structured and they generate better images and videos. -## Prompt engineering +This guide shows you how to write effective prompts and introduces techniques that make them stronger. -> [!TIP] -> This is not an exhaustive guide on prompt engineering, but it will help you understand the necessary parts of a good prompt. We encourage you to continue experimenting with different prompts and combine them in new ways to see what works best. As you write more prompts, you'll develop an intuition for what works and what doesn't! +## Writing good prompts -New diffusion models do a pretty good job of generating high-quality images from a basic prompt, but it is still important to create a well-written prompt to get the best results. Here are a few tips for writing a good prompt: +Every effective prompt needs three core elements. -1. What is the image *medium*? Is it a photo, a painting, a 3D illustration, or something else? -2. What is the image *subject*? Is it a person, animal, object, or scene? -3. What *details* would you like to see in the image? This is where you can get really creative and have a lot of fun experimenting with different words to bring your image to life. For example, what is the lighting like? What is the vibe and aesthetic? What kind of art or illustration style are you looking for? The more specific and precise words you use, the better the model will understand what you want to generate. +1. Subject - what you want to generate. Start your prompt here. +2. Style - the medium or aesthetic. How should it look? +3. Context - details about actions, setting, and mood. + +Use these elements as a structured narrative, not a keyword list. Modern models understand language better than keyword matching. Start simple, then add details. + +Context is especially important for creating better prompts. Try adding lighting, artistic details, and mood.
-
- -
"A photo of a banana-shaped couch in a living room"
+
+ +
A cute cat lounges on a leaf in a pool during a peaceful summer afternoon, in lofi art style, illustration.
-
- -
"A vibrant yellow banana-shaped couch sits in a cozy living room, its curve cradling a pile of colorful cushions. on the wooden floor, a patterned rug adds a touch of eclectic charm, and a potted plant sits in the corner, reaching towards the sunlight filtering through the windows"
+
+ +
A cute cat lounges on a floating leaf in a sparkling pool during a peaceful summer afternoon. Clear reflections ripple across the water, with sunlight casting soft, smooth highlights. The illustration is detailed and polished, with elegant lines and harmonious colors, evoking a relaxing, serene, and whimsical lofi mood, anime-inspired and visually comforting.
-## Prompt enhancing with GPT2 - -Prompt enhancing is a technique for quickly improving prompt quality without spending too much effort constructing one. It uses a model like GPT2 pretrained on Stable Diffusion text prompts to automatically enrich a prompt with additional important keywords to generate high-quality images. - -The technique works by curating a list of specific keywords and forcing the model to generate those words to enhance the original prompt. This way, your prompt can be "a cat" and GPT2 can enhance the prompt to "cinematic film still of a cat basking in the sun on a roof in Turkey, highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain quality sharp focus beautiful detailed intricate stunning amazing epic". +Be specific and add context. Use photography terms like lens type, focal length, camera angles, and depth of field. > [!TIP] -> You should also use a [*offset noise*](https://www.crosslabs.org//blog/diffusion-with-offset-noise) LoRA to improve the contrast in bright and dark images and create better lighting overall. This [LoRA](https://hf.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_offset_example-lora_1.0.safetensors) is available from [stabilityai/stable-diffusion-xl-base-1.0](https://hf.co/stabilityai/stable-diffusion-xl-base-1.0). - -Start by defining certain styles and a list of words (you can check out a more comprehensive list of [words](https://hf.co/LykosAI/GPT-Prompt-Expansion-Fooocus-v2/blob/main/positive.txt) and [styles](https://github.com/lllyasviel/Fooocus/tree/main/sdxl_styles) used by Fooocus) to enhance a prompt with. - -```py -import torch -from transformers import GenerationConfig, GPT2LMHeadModel, GPT2Tokenizer, LogitsProcessor, LogitsProcessorList -from diffusers import StableDiffusionXLPipeline - -styles = { - "cinematic": "cinematic film still of {prompt}, highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain", - "anime": "anime artwork of {prompt}, anime style, key visual, vibrant, studio anime, highly detailed", - "photographic": "cinematic photo of {prompt}, 35mm photograph, film, professional, 4k, highly detailed", - "comic": "comic of {prompt}, graphic illustration, comic art, graphic novel art, vibrant, highly detailed", - "lineart": "line art drawing {prompt}, professional, sleek, modern, minimalist, graphic, line art, vector graphics", - "pixelart": " pixel-art {prompt}, low-res, blocky, pixel art style, 8-bit graphics", -} - -words = [ - "aesthetic", "astonishing", "beautiful", "breathtaking", "composition", "contrasted", "epic", "moody", "enhanced", - "exceptional", "fascinating", "flawless", "glamorous", "glorious", "illumination", "impressive", "improved", - "inspirational", "magnificent", "majestic", "hyperrealistic", "smooth", "sharp", "focus", "stunning", "detailed", - "intricate", "dramatic", "high", "quality", "perfect", "light", "ultra", "highly", "radiant", "satisfying", - "soothing", "sophisticated", "stylish", "sublime", "terrific", "touching", "timeless", "wonderful", "unbelievable", - "elegant", "awesome", "amazing", "dynamic", "trendy", -] -``` - -You may have noticed in the `words` list, there are certain words that can be paired together to create something more meaningful. For example, the words "high" and "quality" can be combined to create "high quality". Let's pair these words together and remove the words that can't be paired. - -```py -word_pairs = ["highly detailed", "high quality", "enhanced quality", "perfect composition", "dynamic light"] - -def find_and_order_pairs(s, pairs): - words = s.split() - found_pairs = [] - for pair in pairs: - pair_words = pair.split() - if pair_words[0] in words and pair_words[1] in words: - found_pairs.append(pair) - words.remove(pair_words[0]) - words.remove(pair_words[1]) - - for word in words[:]: - for pair in pairs: - if word in pair.split(): - words.remove(word) - break - ordered_pairs = ", ".join(found_pairs) - remaining_s = ", ".join(words) - return ordered_pairs, remaining_s -``` - -Next, implement a custom [`~transformers.LogitsProcessor`] class that assigns tokens in the `words` list a value of 0 and assigns tokens not in the `words` list a negative value so they aren't picked during generation. This way, generation is biased towards words in the `words` list. After a word from the list is used, it is also assigned a negative value so it isn't picked again. - -```py -class CustomLogitsProcessor(LogitsProcessor): - def __init__(self, bias): - super().__init__() - self.bias = bias - - def __call__(self, input_ids, scores): - if len(input_ids.shape) == 2: - last_token_id = input_ids[0, -1] - self.bias[last_token_id] = -1e10 - return scores + self.bias - -word_ids = [tokenizer.encode(word, add_prefix_space=True)[0] for word in words] -bias = torch.full((tokenizer.vocab_size,), -float("Inf")).to("cuda") -bias[word_ids] = 0 -processor = CustomLogitsProcessor(bias) -processor_list = LogitsProcessorList([processor]) -``` - -Combine the prompt and the `cinematic` style prompt defined in the `styles` dictionary earlier. - -```py -prompt = "a cat basking in the sun on a roof in Turkey" -style = "cinematic" - -prompt = styles[style].format(prompt=prompt) -prompt -"cinematic film still of a cat basking in the sun on a roof in Turkey, highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain" -``` - -Load a GPT2 tokenizer and model from the [Gustavosta/MagicPrompt-Stable-Diffusion](https://huggingface.co/Gustavosta/MagicPrompt-Stable-Diffusion) checkpoint (this specific checkpoint is trained to generate prompts) to enhance the prompt. - -```py -tokenizer = GPT2Tokenizer.from_pretrained("Gustavosta/MagicPrompt-Stable-Diffusion") -model = GPT2LMHeadModel.from_pretrained("Gustavosta/MagicPrompt-Stable-Diffusion", torch_dtype=torch.float16).to( - "cuda" -) -model.eval() - -inputs = tokenizer(prompt, return_tensors="pt").to("cuda") -token_count = inputs["input_ids"].shape[1] -max_new_tokens = 50 - token_count - -generation_config = GenerationConfig( - penalty_alpha=0.7, - top_k=50, - eos_token_id=model.config.eos_token_id, - pad_token_id=model.config.eos_token_id, - pad_token=model.config.pad_token_id, - do_sample=True, -) - -with torch.no_grad(): - generated_ids = model.generate( - input_ids=inputs["input_ids"], - attention_mask=inputs["attention_mask"], - max_new_tokens=max_new_tokens, - generation_config=generation_config, - logits_processor=proccesor_list, - ) -``` - -Then you can combine the input prompt and the generated prompt. Feel free to take a look at what the generated prompt (`generated_part`) is, the word pairs that were found (`pairs`), and the remaining words (`words`). This is all packed together in the `enhanced_prompt`. - -```py -output_tokens = [tokenizer.decode(generated_id, skip_special_tokens=True) for generated_id in generated_ids] -input_part, generated_part = output_tokens[0][: len(prompt)], output_tokens[0][len(prompt) :] -pairs, words = find_and_order_pairs(generated_part, word_pairs) -formatted_generated_part = pairs + ", " + words -enhanced_prompt = input_part + ", " + formatted_generated_part -enhanced_prompt -["cinematic film still of a cat basking in the sun on a roof in Turkey, highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain quality sharp focus beautiful detailed intricate stunning amazing epic"] -``` - -Finally, load a pipeline and the offset noise LoRA with a *low weight* to generate an image with the enhanced prompt. - -```py -pipeline = StableDiffusionXLPipeline.from_pretrained( - "RunDiffusion/Juggernaut-XL-v9", torch_dtype=torch.float16, variant="fp16" -).to("cuda") - -pipeline.load_lora_weights( - "stabilityai/stable-diffusion-xl-base-1.0", - weight_name="sd_xl_offset_example-lora_1.0.safetensors", - adapter_name="offset", -) -pipeline.set_adapters(["offset"], adapter_weights=[0.2]) - -image = pipeline( - enhanced_prompt, - width=1152, - height=896, - guidance_scale=7.5, - num_inference_steps=25, -).images[0] -image -``` - -
-
- -
"a cat basking in the sun on a roof in Turkey"
-
-
- -
"cinematic film still of a cat basking in the sun on a roof in Turkey, highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain"
-
-
+> Try a [prompt enhancer](https://huggingface.co/models?sort=downloads&search=prompt+enhancer) to help improve your prompt structure. ## Prompt weighting -Prompt weighting provides a way to emphasize or de-emphasize certain parts of a prompt, allowing for more control over the generated image. A prompt can include several concepts, which gets turned into contextualized text embeddings. The embeddings are used by the model to condition its cross-attention layers to generate an image (read the Stable Diffusion [blog post](https://huggingface.co/blog/stable_diffusion) to learn more about how it works). +Prompt weighting makes some words stronger and others weaker. It scales attention scores so you control how much influence each concept has. -Prompt weighting works by increasing or decreasing the scale of the text embedding vector that corresponds to its concept in the prompt because you may not necessarily want the model to focus on all concepts equally. The easiest way to prepare the prompt embeddings is to use [Stable Diffusion Long Prompt Weighted Embedding](https://github.com/xhinker/sd_embed) (sd_embed). Once you have the prompt-weighted embeddings, you can pass them to any pipeline that has a [prompt_embeds](https://huggingface.co/docs/diffusers/en/api/pipelines/stable_diffusion/text2img#diffusers.StableDiffusionPipeline.__call__.prompt_embeds) (and optionally [negative_prompt_embeds](https://huggingface.co/docs/diffusers/en/api/pipelines/stable_diffusion/text2img#diffusers.StableDiffusionPipeline.__call__.negative_prompt_embeds)) parameter, such as [`StableDiffusionPipeline`], [`StableDiffusionControlNetPipeline`], and [`StableDiffusionXLPipeline`]. +Diffusers handles this through `prompt_embeds` and `pooled_prompt_embeds` arguments which take scaled text embedding vectors. Use the [sd_embed](https://github.com/xhinker/sd_embed) library to generate these embeddings. It also supports longer prompts. -> [!TIP] -> If your favorite pipeline doesn't have a `prompt_embeds` parameter, please open an [issue](https://github.com/huggingface/diffusers/issues/new/choose) so we can add it! - -This guide will show you how to weight your prompts with sd_embed. - -Before you begin, make sure you have the latest version of sd_embed installed: - -```bash -pip install git+https://github.com/xhinker/sd_embed.git@main -``` - -For this example, let's use [`StableDiffusionXLPipeline`]. +> [!NOTE] +> The sd_embed library only supports Stable Diffusion, Stable Diffusion XL, Stable Diffusion 3, Stable Cascade, and Flux. Prompt weighting doesn't necessarily help for newer models like Flux which already has very good prompt adherence. ```py -from diffusers import StableDiffusionXLPipeline, UniPCMultistepScheduler -import torch - -pipe = StableDiffusionXLPipeline.from_pretrained("Lykon/dreamshaper-xl-1-0", torch_dtype=torch.float16) -pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) -pipe.to("cuda") +!uv pip install git+https://github.com/xhinker/sd_embed.git@main ``` -To upweight or downweight a concept, surround the text with parentheses. More parentheses applies a heavier weight on the text. You can also append a numerical multiplier to the text to indicate how much you want to increase or decrease its weights by. +Format weighted text with numerical multipliers or parentheses. More parentheses mean stronger weighting. | format | multiplier | |---|---| -| `(hippo)` | increase by 1.1x | -| `((hippo))` | increase by 1.21x | -| `(hippo:1.5)` | increase by 1.5x | -| `(hippo:0.5)` | decrease by 4x | +| `(cat)` | increase by 1.1x | +| `((cat))` | increase by 1.21x | +| `(cat:1.5)` | increase by 1.5x | +| `(cat:0.5)` | decrease by 4x | -Create a prompt and use a combination of parentheses and numerical multipliers to upweight various text. +Create a weighted prompt and pass it to [get_weighted_text_embeddings_sdxl](https://github.com/xhinker/sd_embed/blob/4a47f71150a22942fa606fb741a1c971d95ba56f/src/sd_embed/embedding_funcs.py#L405) to generate embeddings. + +> [!TIP] +> You could also pass negative prompts to `negative_prompt_embeds` and `negative_pooled_prompt_embeds`. ```py +import torch +from diffusers import DiffusionPipeline from sd_embed.embedding_funcs import get_weighted_text_embeddings_sdxl -prompt = """A whimsical and creative image depicting a hybrid creature that is a mix of a waffle and a hippopotamus. -This imaginative creature features the distinctive, bulky body of a hippo, -but with a texture and appearance resembling a golden-brown, crispy waffle. -The creature might have elements like waffle squares across its skin and a syrup-like sheen. -It's set in a surreal environment that playfully combines a natural water habitat of a hippo with elements of a breakfast table setting, -possibly including oversized utensils or plates in the background. -The image should evoke a sense of playful absurdity and culinary fantasy. -""" - -neg_prompt = """\ -skin spots,acnes,skin blemishes,age spot,(ugly:1.2),(duplicate:1.2),(morbid:1.21),(mutilated:1.2),\ -(tranny:1.2),mutated hands,(poorly drawn hands:1.5),blurry,(bad anatomy:1.2),(bad proportions:1.3),\ -extra limbs,(disfigured:1.2),(missing arms:1.2),(extra legs:1.2),(fused fingers:1.5),\ -(too many fingers:1.5),(unclear eyes:1.2),lowers,bad hands,missing fingers,extra digit,\ -bad hands,missing fingers,(extra arms and legs),(worst quality:2),(low quality:2),\ -(normal quality:2),lowres,((monochrome)),((grayscale)) -""" -``` - -Use the `get_weighted_text_embeddings_sdxl` function to generate the prompt embeddings and the negative prompt embeddings. It'll also generated the pooled and negative pooled prompt embeddings since you're using the SDXL model. - -> [!TIP] -> You can safely ignore the error message below about the token index length exceeding the models maximum sequence length. All your tokens will be used in the embedding process. -> -> ``` -> Token indices sequence length is longer than the specified maximum sequence length for this model -> ``` - -```py -( - prompt_embeds, - prompt_neg_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds -) = get_weighted_text_embeddings_sdxl( - pipe, - prompt=prompt, - neg_prompt=neg_prompt +pipeline = DiffusionPipeline.from_pretrained( + "Lykon/dreamshaper-xl-1-0", torch_dtype=torch.bfloat16, device_map="cuda" ) -image = pipe( - prompt_embeds=prompt_embeds, - negative_prompt_embeds=prompt_neg_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - num_inference_steps=30, - height=1024, - width=1024 + 512, - guidance_scale=4.0, - generator=torch.Generator("cuda").manual_seed(2) -).images[0] -image +prompt = """ +A (cute cat:1.4) lounges on a (floating leaf:1.2) in a (sparkling pool:1.1) during a peaceful summer afternoon. +Gentle ripples reflect pastel skies, while (sunlight:1.1) casts soft highlights. The illustration is smooth and polished +with elegant, sketchy lines and subtle gradients, evoking a ((whimsical, nostalgic, dreamy lofi atmosphere:2.0)), +(anime-inspired:1.6), calming, comforting, and visually serene. +""" + +prompt_embeds, _, pooled_prompt_embeds, *_ = get_weighted_text_embeddings_sdxl(pipeline, prompt=prompt) +``` + +Pass the embeddings to `prompt_embeds` and `pooled_prompt_embeds` to generate your image. + +```py +image = pipeline(prompt_embeds=prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds).images[0] ```
- +
-> [!TIP] -> Refer to the [sd_embed](https://github.com/xhinker/sd_embed) repository for additional details about long prompt weighting for FLUX.1, Stable Cascade, and Stable Diffusion 1.5. - -### Textual inversion - -[Textual inversion](../training/text_inversion) is a technique for learning a specific concept from some images which you can use to generate new images conditioned on that concept. - -Create a pipeline and use the [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] function to load the textual inversion embeddings (feel free to browse the [Stable Diffusion Conceptualizer](https://huggingface.co/spaces/sd-concepts-library/stable-diffusion-conceptualizer) for 100+ trained concepts): - -```py -import torch -from diffusers import StableDiffusionPipeline - -pipe = StableDiffusionPipeline.from_pretrained( - "stable-diffusion-v1-5/stable-diffusion-v1-5", - torch_dtype=torch.float16, -).to("cuda") -pipe.load_textual_inversion("sd-concepts-library/midjourney-style") -``` - -Add the `` text to the prompt to trigger the textual inversion. - -```py -from sd_embed.embedding_funcs import get_weighted_text_embeddings_sd15 - -prompt = """ A whimsical and creative image depicting a hybrid creature that is a mix of a waffle and a hippopotamus. -This imaginative creature features the distinctive, bulky body of a hippo, -but with a texture and appearance resembling a golden-brown, crispy waffle. -The creature might have elements like waffle squares across its skin and a syrup-like sheen. -It's set in a surreal environment that playfully combines a natural water habitat of a hippo with elements of a breakfast table setting, -possibly including oversized utensils or plates in the background. -The image should evoke a sense of playful absurdity and culinary fantasy. -""" - -neg_prompt = """\ -skin spots,acnes,skin blemishes,age spot,(ugly:1.2),(duplicate:1.2),(morbid:1.21),(mutilated:1.2),\ -(tranny:1.2),mutated hands,(poorly drawn hands:1.5),blurry,(bad anatomy:1.2),(bad proportions:1.3),\ -extra limbs,(disfigured:1.2),(missing arms:1.2),(extra legs:1.2),(fused fingers:1.5),\ -(too many fingers:1.5),(unclear eyes:1.2),lowers,bad hands,missing fingers,extra digit,\ -bad hands,missing fingers,(extra arms and legs),(worst quality:2),(low quality:2),\ -(normal quality:2),lowres,((monochrome)),((grayscale)) -""" -``` - -Use the `get_weighted_text_embeddings_sd15` function to generate the prompt embeddings and the negative prompt embeddings. - -```py -( - prompt_embeds, - prompt_neg_embeds, -) = get_weighted_text_embeddings_sd15( - pipe, - prompt=prompt, - neg_prompt=neg_prompt -) - -image = pipe( - prompt_embeds=prompt_embeds, - negative_prompt_embeds=prompt_neg_embeds, - height=768, - width=896, - guidance_scale=4.0, - generator=torch.Generator("cuda").manual_seed(2) -).images[0] -image -``` - -
- -
- -### DreamBooth - -[DreamBooth](../training/dreambooth) is a technique for generating contextualized images of a subject given just a few images of the subject to train on. It is similar to textual inversion, but DreamBooth trains the full model whereas textual inversion only fine-tunes the text embeddings. This means you should use [`~DiffusionPipeline.from_pretrained`] to load the DreamBooth model (feel free to browse the [Stable Diffusion Dreambooth Concepts Library](https://huggingface.co/sd-dreambooth-library) for 100+ trained models): - -```py -import torch -from diffusers import DiffusionPipeline, UniPCMultistepScheduler - -pipe = DiffusionPipeline.from_pretrained("sd-dreambooth-library/dndcoverart-v1", torch_dtype=torch.float16).to("cuda") -pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) -``` - -Depending on the model you use, you'll need to incorporate the model's unique identifier into your prompt. For example, the `dndcoverart-v1` model uses the identifier `dndcoverart`: - -```py -from sd_embed.embedding_funcs import get_weighted_text_embeddings_sd15 - -prompt = """dndcoverart of A whimsical and creative image depicting a hybrid creature that is a mix of a waffle and a hippopotamus. -This imaginative creature features the distinctive, bulky body of a hippo, -but with a texture and appearance resembling a golden-brown, crispy waffle. -The creature might have elements like waffle squares across its skin and a syrup-like sheen. -It's set in a surreal environment that playfully combines a natural water habitat of a hippo with elements of a breakfast table setting, -possibly including oversized utensils or plates in the background. -The image should evoke a sense of playful absurdity and culinary fantasy. -""" - -neg_prompt = """\ -skin spots,acnes,skin blemishes,age spot,(ugly:1.2),(duplicate:1.2),(morbid:1.21),(mutilated:1.2),\ -(tranny:1.2),mutated hands,(poorly drawn hands:1.5),blurry,(bad anatomy:1.2),(bad proportions:1.3),\ -extra limbs,(disfigured:1.2),(missing arms:1.2),(extra legs:1.2),(fused fingers:1.5),\ -(too many fingers:1.5),(unclear eyes:1.2),lowers,bad hands,missing fingers,extra digit,\ -bad hands,missing fingers,(extra arms and legs),(worst quality:2),(low quality:2),\ -(normal quality:2),lowres,((monochrome)),((grayscale)) -""" - -( - prompt_embeds - , prompt_neg_embeds -) = get_weighted_text_embeddings_sd15( - pipe - , prompt = prompt - , neg_prompt = neg_prompt -) -``` - -
- -
+Prompt weighting works with [Textual inversion](./textual_inversion_inference) and [DreamBooth](./dreambooth) adapters too. \ No newline at end of file From b4e6dc3037e75e2dc16466b914b1be597b06be9a Mon Sep 17 00:00:00 2001 From: Steven Liu <59462357+stevhliu@users.noreply.github.com> Date: Tue, 14 Oct 2025 18:12:10 -0700 Subject: [PATCH 065/123] [docs] Fix broken links (#12487) fix broken links --- docs/source/en/api/loaders/ip_adapter.md | 2 +- docs/source/en/api/loaders/lora.md | 2 +- docs/source/en/api/loaders/peft.md | 2 +- .../en/api/loaders/textual_inversion.md | 2 +- docs/source/en/api/loaders/transformer_sd3.md | 2 +- docs/source/en/api/loaders/unet.md | 2 +- docs/source/en/api/pipelines/flux.md | 2 +- docs/source/en/api/pipelines/hidream.md | 2 +- docs/source/en/api/pipelines/overview.md | 2 +- .../stable_diffusion/stable_diffusion_3.md | 2 +- .../en/api/pipelines/stable_diffusion/svd.md | 2 +- docs/source/en/api/pipelines/text_to_video.md | 2 +- docs/source/en/training/dreambooth.md | 2 +- docs/source/en/training/lcm_distill.md | 4 +-- docs/source/en/training/lora.md | 2 +- docs/source/en/training/text2image.md | 2 +- docs/source/en/training/text_inversion.md | 3 +- .../using-diffusers/controlling_generation.md | 34 ------------------- .../en/using-diffusers/inference_with_lcm.md | 2 +- .../inference_with_tcd_lora.md | 4 +-- docs/source/en/using-diffusers/sdxl.md | 2 +- .../en/using-diffusers/write_own_pipeline.md | 2 +- 22 files changed, 23 insertions(+), 58 deletions(-) diff --git a/docs/source/en/api/loaders/ip_adapter.md b/docs/source/en/api/loaders/ip_adapter.md index c2c45bee10..508e6d4ee6 100644 --- a/docs/source/en/api/loaders/ip_adapter.md +++ b/docs/source/en/api/loaders/ip_adapter.md @@ -15,7 +15,7 @@ specific language governing permissions and limitations under the License. [IP-Adapter](https://hf.co/papers/2308.06721) is a lightweight adapter that enables prompting a diffusion model with an image. This method decouples the cross-attention layers of the image and text features. The image features are generated from an image encoder. > [!TIP] -> Learn how to load an IP-Adapter checkpoint and image in the IP-Adapter [loading](../../using-diffusers/loading_adapters#ip-adapter) guide, and you can see how to use it in the [usage](../../using-diffusers/ip_adapter) guide. +> Learn how to load and use an IP-Adapter checkpoint and image in the [IP-Adapter](../../using-diffusers/ip_adapter) guide,. ## IPAdapterMixin diff --git a/docs/source/en/api/loaders/lora.md b/docs/source/en/api/loaders/lora.md index bf22a32d74..b1d1ffb634 100644 --- a/docs/source/en/api/loaders/lora.md +++ b/docs/source/en/api/loaders/lora.md @@ -34,7 +34,7 @@ LoRA is a fast and lightweight training method that inserts and trains a signifi - [`LoraBaseMixin`] provides a base class with several utility methods to fuse, unfuse, unload, LoRAs and more. > [!TIP] -> To learn more about how to load LoRA weights, see the [LoRA](../../using-diffusers/loading_adapters#lora) loading guide. +> To learn more about how to load LoRA weights, see the [LoRA](../../tutorials/using_peft_for_inference) loading guide. ## LoraBaseMixin diff --git a/docs/source/en/api/loaders/peft.md b/docs/source/en/api/loaders/peft.md index 5508509c88..c514766dd8 100644 --- a/docs/source/en/api/loaders/peft.md +++ b/docs/source/en/api/loaders/peft.md @@ -12,7 +12,7 @@ specific language governing permissions and limitations under the License. # PEFT -Diffusers supports loading adapters such as [LoRA](../../using-diffusers/loading_adapters) with the [PEFT](https://huggingface.co/docs/peft/index) library with the [`~loaders.peft.PeftAdapterMixin`] class. This allows modeling classes in Diffusers like [`UNet2DConditionModel`], [`SD3Transformer2DModel`] to operate with an adapter. +Diffusers supports loading adapters such as [LoRA](../../tutorials/using_peft_for_inference) with the [PEFT](https://huggingface.co/docs/peft/index) library with the [`~loaders.peft.PeftAdapterMixin`] class. This allows modeling classes in Diffusers like [`UNet2DConditionModel`], [`SD3Transformer2DModel`] to operate with an adapter. > [!TIP] > Refer to the [Inference with PEFT](../../tutorials/using_peft_for_inference.md) tutorial for an overview of how to use PEFT in Diffusers for inference. diff --git a/docs/source/en/api/loaders/textual_inversion.md b/docs/source/en/api/loaders/textual_inversion.md index 2cb54ce4ea..5e8bfac255 100644 --- a/docs/source/en/api/loaders/textual_inversion.md +++ b/docs/source/en/api/loaders/textual_inversion.md @@ -17,7 +17,7 @@ Textual Inversion is a training method for personalizing models by learning new [`TextualInversionLoaderMixin`] provides a function for loading Textual Inversion embeddings from Diffusers and Automatic1111 into the text encoder and loading a special token to activate the embeddings. > [!TIP] -> To learn more about how to load Textual Inversion embeddings, see the [Textual Inversion](../../using-diffusers/loading_adapters#textual-inversion) loading guide. +> To learn more about how to load Textual Inversion embeddings, see the [Textual Inversion](../../using-diffusers/textual_inversion_inference) loading guide. ## TextualInversionLoaderMixin diff --git a/docs/source/en/api/loaders/transformer_sd3.md b/docs/source/en/api/loaders/transformer_sd3.md index cc3ec0da14..2c8b81b59c 100644 --- a/docs/source/en/api/loaders/transformer_sd3.md +++ b/docs/source/en/api/loaders/transformer_sd3.md @@ -17,7 +17,7 @@ This class is useful when *only* loading weights into a [`SD3Transformer2DModel` The [`SD3Transformer2DLoadersMixin`] class currently only loads IP-Adapter weights, but will be used in the future to save weights and load LoRAs. > [!TIP] -> To learn more about how to load LoRA weights, see the [LoRA](../../using-diffusers/loading_adapters#lora) loading guide. +> To learn more about how to load LoRA weights, see the [LoRA](../../tutorials/using_peft_for_inference) loading guide. ## SD3Transformer2DLoadersMixin diff --git a/docs/source/en/api/loaders/unet.md b/docs/source/en/api/loaders/unet.md index 7450e03e58..50d210bbf5 100644 --- a/docs/source/en/api/loaders/unet.md +++ b/docs/source/en/api/loaders/unet.md @@ -17,7 +17,7 @@ Some training methods - like LoRA and Custom Diffusion - typically target the UN The [`UNet2DConditionLoadersMixin`] class provides functions for loading and saving weights, fusing and unfusing LoRAs, disabling and enabling LoRAs, and setting and deleting adapters. > [!TIP] -> To learn more about how to load LoRA weights, see the [LoRA](../../using-diffusers/loading_adapters#lora) loading guide. +> To learn more about how to load LoRA weights, see the [LoRA](../../tutorials/using_peft_for_inference) guide. ## UNet2DConditionLoadersMixin diff --git a/docs/source/en/api/pipelines/flux.md b/docs/source/en/api/pipelines/flux.md index 1a89de98e4..358b8139c7 100644 --- a/docs/source/en/api/pipelines/flux.md +++ b/docs/source/en/api/pipelines/flux.md @@ -418,7 +418,7 @@ When unloading the Control LoRA weights, call `pipe.unload_lora_weights(reset_to ## IP-Adapter > [!TIP] -> Check out [IP-Adapter](../../../using-diffusers/ip_adapter) to learn more about how IP-Adapters work. +> Check out [IP-Adapter](../../using-diffusers/ip_adapter) to learn more about how IP-Adapters work. An IP-Adapter lets you prompt Flux with images, in addition to the text prompt. This is especially useful when describing complex concepts that are difficult to articulate through text alone and you have reference images. diff --git a/docs/source/en/api/pipelines/hidream.md b/docs/source/en/api/pipelines/hidream.md index acfcef93e0..add4ad3132 100644 --- a/docs/source/en/api/pipelines/hidream.md +++ b/docs/source/en/api/pipelines/hidream.md @@ -21,7 +21,7 @@ ## Available models -The following models are available for the [`HiDreamImagePipeline`](text-to-image) pipeline: +The following models are available for the [`HiDreamImagePipeline`] pipeline: | Model name | Description | |:---|:---| diff --git a/docs/source/en/api/pipelines/overview.md b/docs/source/en/api/pipelines/overview.md index ce883931df..22fcf560ea 100644 --- a/docs/source/en/api/pipelines/overview.md +++ b/docs/source/en/api/pipelines/overview.md @@ -32,7 +32,7 @@ The table below lists all the pipelines currently available in 🤗 Diffusers an | [Attend-and-Excite](attend_and_excite) | text2image | | [AudioLDM](audioldm) | text2audio | | [AudioLDM2](audioldm2) | text2audio | -| [AuraFlow](auraflow) | text2image | +| [AuraFlow](aura_flow) | text2image | | [BLIP Diffusion](blip_diffusion) | text2image | | [Bria 3.2](bria_3_2) | text2image | | [CogVideoX](cogvideox) | text2video | diff --git a/docs/source/en/api/pipelines/stable_diffusion/stable_diffusion_3.md b/docs/source/en/api/pipelines/stable_diffusion/stable_diffusion_3.md index 3c49df101c..4c7e5b1073 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/stable_diffusion_3.md +++ b/docs/source/en/api/pipelines/stable_diffusion/stable_diffusion_3.md @@ -271,7 +271,7 @@ Check out the full script [here](https://gist.github.com/sayakpaul/508d89d7aad4f Quantization helps reduce the memory requirements of very large models by storing model weights in a lower precision data type. However, quantization may have varying impact on video quality depending on the video model. -Refer to the [Quantization](../../quantization/overview) overview to learn more about supported quantization backends and selecting a quantization backend that supports your use case. The example below demonstrates how to load a quantized [`StableDiffusion3Pipeline`] for inference with bitsandbytes. +Refer to the [Quantization](../../../quantization/overview) overview to learn more about supported quantization backends and selecting a quantization backend that supports your use case. The example below demonstrates how to load a quantized [`StableDiffusion3Pipeline`] for inference with bitsandbytes. ```py import torch diff --git a/docs/source/en/api/pipelines/stable_diffusion/svd.md b/docs/source/en/api/pipelines/stable_diffusion/svd.md index 0c33c06007..a00dd3ef6d 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/svd.md +++ b/docs/source/en/api/pipelines/stable_diffusion/svd.md @@ -29,7 +29,7 @@ The abstract from the paper is: Video generation is memory-intensive and one way to reduce your memory usage is to set `enable_forward_chunking` on the pipeline's UNet so you don't run the entire feedforward layer at once. Breaking it up into chunks in a loop is more efficient. -Check out the [Text or image-to-video](text-img2vid) guide for more details about how certain parameters can affect video generation and how to optimize inference by reducing memory usage. +Check out the [Text or image-to-video](../../../using-diffusers/text-img2vid) guide for more details about how certain parameters can affect video generation and how to optimize inference by reducing memory usage. ## StableVideoDiffusionPipeline diff --git a/docs/source/en/api/pipelines/text_to_video.md b/docs/source/en/api/pipelines/text_to_video.md index d7c37be637..d9f6d8e722 100644 --- a/docs/source/en/api/pipelines/text_to_video.md +++ b/docs/source/en/api/pipelines/text_to_video.md @@ -172,7 +172,7 @@ Here are some sample outputs: Video generation is memory-intensive and one way to reduce your memory usage is to set `enable_forward_chunking` on the pipeline's UNet so you don't run the entire feedforward layer at once. Breaking it up into chunks in a loop is more efficient. -Check out the [Text or image-to-video](text-img2vid) guide for more details about how certain parameters can affect video generation and how to optimize inference by reducing memory usage. +Check out the [Text or image-to-video](../../using-diffusers/text-img2vid) guide for more details about how certain parameters can affect video generation and how to optimize inference by reducing memory usage. > [!TIP] > Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. diff --git a/docs/source/en/training/dreambooth.md b/docs/source/en/training/dreambooth.md index 81ed09c9d0..2302739a0e 100644 --- a/docs/source/en/training/dreambooth.md +++ b/docs/source/en/training/dreambooth.md @@ -548,4 +548,4 @@ Training the DeepFloyd IF model can be challenging, but here are some tips that Congratulations on training your DreamBooth model! To learn more about how to use your new model, the following guide may be helpful: -- Learn how to [load a DreamBooth](../using-diffusers/loading_adapters) model for inference if you trained your model with LoRA. \ No newline at end of file +- Learn how to [load a DreamBooth](../using-diffusers/dreambooth) model for inference if you trained your model with LoRA. \ No newline at end of file diff --git a/docs/source/en/training/lcm_distill.md b/docs/source/en/training/lcm_distill.md index 232f2eceed..4750f15036 100644 --- a/docs/source/en/training/lcm_distill.md +++ b/docs/source/en/training/lcm_distill.md @@ -75,7 +75,7 @@ accelerate launch train_lcm_distill_sd_wds.py \ Most of the parameters are identical to the parameters in the [Text-to-image](text2image#script-parameters) training guide, so you'll focus on the parameters that are relevant to latent consistency distillation in this guide. - `--pretrained_teacher_model`: the path to a pretrained latent diffusion model to use as the teacher model -- `--pretrained_vae_model_name_or_path`: path to a pretrained VAE; the SDXL VAE is known to suffer from numerical instability, so this parameter allows you to specify an alternative VAE (like this [VAE]((https://huggingface.co/madebyollin/sdxl-vae-fp16-fix)) by madebyollin which works in fp16) +- `--pretrained_vae_model_name_or_path`: path to a pretrained VAE; the SDXL VAE is known to suffer from numerical instability, so this parameter allows you to specify an alternative VAE (like this [VAE](https://huggingface.co/madebyollin/sdxl-vae-fp16-fix)) by madebyollin which works in fp16) - `--w_min` and `--w_max`: the minimum and maximum guidance scale values for guidance scale sampling - `--num_ddim_timesteps`: the number of timesteps for DDIM sampling - `--loss_type`: the type of loss (L2 or Huber) to calculate for latent consistency distillation; Huber loss is generally preferred because it's more robust to outliers @@ -245,5 +245,5 @@ The SDXL training script is discussed in more detail in the [SDXL training](sdxl Congratulations on distilling a LCM model! To learn more about LCM, the following may be helpful: -- Learn how to use [LCMs for inference](../using-diffusers/lcm) for text-to-image, image-to-image, and with LoRA checkpoints. +- Learn how to use [LCMs for inference](../using-diffusers/inference_with_lcm) for text-to-image, image-to-image, and with LoRA checkpoints. - Read the [SDXL in 4 steps with Latent Consistency LoRAs](https://huggingface.co/blog/lcm_lora) blog post to learn more about SDXL LCM-LoRA's for super fast inference, quality comparisons, benchmarks, and more. diff --git a/docs/source/en/training/lora.md b/docs/source/en/training/lora.md index 45a234b76a..efb170e329 100644 --- a/docs/source/en/training/lora.md +++ b/docs/source/en/training/lora.md @@ -198,5 +198,5 @@ image = pipeline("A naruto with blue eyes").images[0] Congratulations on training a new model with LoRA! To learn more about how to use your new model, the following guides may be helpful: -- Learn how to [load different LoRA formats](../using-diffusers/loading_adapters#LoRA) trained using community trainers like Kohya and TheLastBen. +- Learn how to [load different LoRA formats](../tutorials/using_peft_for_inference) trained using community trainers like Kohya and TheLastBen. - Learn how to use and [combine multiple LoRA's](../tutorials/using_peft_for_inference) with PEFT for inference. diff --git a/docs/source/en/training/text2image.md b/docs/source/en/training/text2image.md index a9327457c7..d11e55e910 100644 --- a/docs/source/en/training/text2image.md +++ b/docs/source/en/training/text2image.md @@ -178,5 +178,5 @@ image.save("yoda-naruto.png") Congratulations on training your own text-to-image model! To learn more about how to use your new model, the following guides may be helpful: -- Learn how to [load LoRA weights](../using-diffusers/loading_adapters#LoRA) for inference if you trained your model with LoRA. +- Learn how to [load LoRA weights](../tutorials/using_peft_for_inference) for inference if you trained your model with LoRA. - Learn more about how certain parameters like guidance scale or techniques such as prompt weighting can help you control inference in the [Text-to-image](../using-diffusers/conditional_image_generation) task guide. diff --git a/docs/source/en/training/text_inversion.md b/docs/source/en/training/text_inversion.md index 0b540107e9..4912b6730a 100644 --- a/docs/source/en/training/text_inversion.md +++ b/docs/source/en/training/text_inversion.md @@ -203,5 +203,4 @@ image.save("cat-train.png") Congratulations on training your own Textual Inversion model! 🎉 To learn more about how to use your new model, the following guides may be helpful: -- Learn how to [load Textual Inversion embeddings](../using-diffusers/loading_adapters) and also use them as negative embeddings. -- Learn how to use [Textual Inversion](textual_inversion_inference) for inference with Stable Diffusion 1/2 and Stable Diffusion XL. \ No newline at end of file +- Learn how to [load Textual Inversion embeddings](../using-diffusers/textual_inversion_inference) and also use them as negative embeddings. \ No newline at end of file diff --git a/docs/source/en/using-diffusers/controlling_generation.md b/docs/source/en/using-diffusers/controlling_generation.md index aed3a8b729..b7b0ea4919 100644 --- a/docs/source/en/using-diffusers/controlling_generation.md +++ b/docs/source/en/using-diffusers/controlling_generation.md @@ -70,32 +70,6 @@ For convenience, we provide a table to denote which methods are inference-only a [InstructPix2Pix](../api/pipelines/pix2pix) is fine-tuned from Stable Diffusion to support editing input images. It takes as inputs an image and a prompt describing an edit, and it outputs the edited image. InstructPix2Pix has been explicitly trained to work well with [InstructGPT](https://openai.com/blog/instruction-following/)-like prompts. -## Pix2Pix Zero - -[Paper](https://huggingface.co/papers/2302.03027) - -[Pix2Pix Zero](../api/pipelines/pix2pix_zero) allows modifying an image so that one concept or subject is translated to another one while preserving general image semantics. - -The denoising process is guided from one conceptual embedding towards another conceptual embedding. The intermediate latents are optimized during the denoising process to push the attention maps towards reference attention maps. The reference attention maps are from the denoising process of the input image and are used to encourage semantic preservation. - -Pix2Pix Zero can be used both to edit synthetic images as well as real images. - -- To edit synthetic images, one first generates an image given a caption. - Next, we generate image captions for the concept that shall be edited and for the new target concept. We can use a model like [Flan-T5](https://huggingface.co/docs/transformers/model_doc/flan-t5) for this purpose. Then, "mean" prompt embeddings for both the source and target concepts are created via the text encoder. Finally, the pix2pix-zero algorithm is used to edit the synthetic image. -- To edit a real image, one first generates an image caption using a model like [BLIP](https://huggingface.co/docs/transformers/model_doc/blip). Then one applies DDIM inversion on the prompt and image to generate "inverse" latents. Similar to before, "mean" prompt embeddings for both source and target concepts are created and finally the pix2pix-zero algorithm in combination with the "inverse" latents is used to edit the image. - -> [!TIP] -> Pix2Pix Zero is the first model that allows "zero-shot" image editing. This means that the model -> can edit an image in less than a minute on a consumer GPU as shown [here](../api/pipelines/pix2pix_zero#usage-example). - -As mentioned above, Pix2Pix Zero includes optimizing the latents (and not any of the UNet, VAE, or the text encoder) to steer the generation toward a specific concept. This means that the overall -pipeline might require more memory than a standard [StableDiffusionPipeline](../api/pipelines/stable_diffusion/text2img). - -> [!TIP] -> An important distinction between methods like InstructPix2Pix and Pix2Pix Zero is that the former -> involves fine-tuning the pre-trained weights while the latter does not. This means that you can -> apply Pix2Pix Zero to any of the available Stable Diffusion models. - ## Attend and Excite [Paper](https://huggingface.co/papers/2301.13826) @@ -178,14 +152,6 @@ multi-concept training by design. Like DreamBooth and Textual Inversion, Custom teach a pre-trained text-to-image diffusion model about new concepts to generate outputs involving the concept(s) of interest. -## Model Editing - -[Paper](https://huggingface.co/papers/2303.08084) - -The [text-to-image model editing pipeline](../api/pipelines/model_editing) helps you mitigate some of the incorrect implicit assumptions a pre-trained text-to-image -diffusion model might make about the subjects present in the input prompt. For example, if you prompt Stable Diffusion to generate images for "A pack of roses", the roses in the generated images -are more likely to be red. This pipeline helps you change that assumption. - ## DiffEdit [Paper](https://huggingface.co/papers/2210.11427) diff --git a/docs/source/en/using-diffusers/inference_with_lcm.md b/docs/source/en/using-diffusers/inference_with_lcm.md index d0a47449ad..cde4168d38 100644 --- a/docs/source/en/using-diffusers/inference_with_lcm.md +++ b/docs/source/en/using-diffusers/inference_with_lcm.md @@ -257,7 +257,7 @@ LCMs are compatible with adapters like LoRA, ControlNet, T2I-Adapter, and Animat ### LoRA -[LoRA](../using-diffusers/loading_adapters#lora) adapters can be rapidly finetuned to learn a new style from just a few images and plugged into a pretrained model to generate images in that style. +[LoRA](../tutorials/using_peft_for_inference) adapters can be rapidly finetuned to learn a new style from just a few images and plugged into a pretrained model to generate images in that style. diff --git a/docs/source/en/using-diffusers/inference_with_tcd_lora.md b/docs/source/en/using-diffusers/inference_with_tcd_lora.md index a4de12b5e7..2aaf9c8aa8 100644 --- a/docs/source/en/using-diffusers/inference_with_tcd_lora.md +++ b/docs/source/en/using-diffusers/inference_with_tcd_lora.md @@ -18,7 +18,7 @@ Trajectory Consistency Distillation (TCD) enables a model to generate higher qua The major advantages of TCD are: -- Better than Teacher: TCD demonstrates superior generative quality at both small and large inference steps and exceeds the performance of [DPM-Solver++(2S)](../../api/schedulers/multistep_dpm_solver) with Stable Diffusion XL (SDXL). There is no additional discriminator or LPIPS supervision included during TCD training. +- Better than Teacher: TCD demonstrates superior generative quality at both small and large inference steps and exceeds the performance of [DPM-Solver++(2S)](../api/schedulers/multistep_dpm_solver) with Stable Diffusion XL (SDXL). There is no additional discriminator or LPIPS supervision included during TCD training. - Flexible Inference Steps: The inference steps for TCD sampling can be freely adjusted without adversely affecting the image quality. @@ -166,7 +166,7 @@ image = pipe( TCD-LoRA also supports other LoRAs trained on different styles. For example, let's load the [TheLastBen/Papercut_SDXL](https://huggingface.co/TheLastBen/Papercut_SDXL) LoRA and fuse it with the TCD-LoRA with the [`~loaders.UNet2DConditionLoadersMixin.set_adapters`] method. > [!TIP] -> Check out the [Merge LoRAs](merge_loras) guide to learn more about efficient merging methods. +> Check out the [Merge LoRAs](../tutorials/using_peft_for_inference#merge) guide to learn more about efficient merging methods. ```python import torch diff --git a/docs/source/en/using-diffusers/sdxl.md b/docs/source/en/using-diffusers/sdxl.md index 79625e0c4a..275394a03c 100644 --- a/docs/source/en/using-diffusers/sdxl.md +++ b/docs/source/en/using-diffusers/sdxl.md @@ -280,7 +280,7 @@ refiner = DiffusionPipeline.from_pretrained( ``` > [!WARNING] -> You can use SDXL refiner with a different base model. For example, you can use the [Hunyuan-DiT](../../api/pipelines/hunyuandit) or [PixArt-Sigma](../../api/pipelines/pixart_sigma) pipelines to generate images with better prompt adherence. Once you have generated an image, you can pass it to the SDXL refiner model to enhance final generation quality. +> You can use SDXL refiner with a different base model. For example, you can use the [Hunyuan-DiT](../api/pipelines/hunyuandit) or [PixArt-Sigma](../api/pipelines/pixart_sigma) pipelines to generate images with better prompt adherence. Once you have generated an image, you can pass it to the SDXL refiner model to enhance final generation quality. Generate an image from the base model, and set the model output to **latent** space: diff --git a/docs/source/en/using-diffusers/write_own_pipeline.md b/docs/source/en/using-diffusers/write_own_pipeline.md index 930b0fe21f..e34727b5da 100644 --- a/docs/source/en/using-diffusers/write_own_pipeline.md +++ b/docs/source/en/using-diffusers/write_own_pipeline.md @@ -280,5 +280,5 @@ This is really what 🧨 Diffusers is designed for: to make it intuitive and eas For your next steps, feel free to: -* Learn how to [build and contribute a pipeline](../using-diffusers/contribute_pipeline) to 🧨 Diffusers. We can't wait and see what you'll come up with! +* Learn how to [build and contribute a pipeline](../conceptual/contribution) to 🧨 Diffusers. We can't wait and see what you'll come up with! * Explore [existing pipelines](../api/pipelines/overview) in the library, and see if you can deconstruct and build a pipeline from scratch using the models and schedulers separately. From 53a10518b9a5ac998d9ed40ae3f3edcaa4eadd89 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Wed, 15 Oct 2025 09:51:18 +0530 Subject: [PATCH 066/123] remove unneeded checkpoint imports. (#12488) --- .../models/autoencoders/autoencoder_kl_hunyuan_video.py | 1 - src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py | 1 - src/diffusers/models/autoencoders/autoencoder_kl_wan.py | 1 - src/diffusers/models/controlnets/controlnet_xs.py | 1 - src/diffusers/models/transformers/stable_audio_transformer.py | 1 - src/diffusers/models/unets/unet_2d_condition.py | 1 - src/diffusers/models/unets/unet_3d_condition.py | 1 - src/diffusers/models/unets/unet_i2vgen_xl.py | 1 - src/diffusers/models/unets/unet_kandinsky3.py | 1 - src/diffusers/models/unets/unet_motion_model.py | 1 - src/diffusers/pipelines/audioldm2/modeling_audioldm2.py | 1 - src/diffusers/pipelines/blip_diffusion/modeling_blip2.py | 1 - .../pipeline_versatile_diffusion_image_variation.py | 1 - .../pipeline_versatile_diffusion_text_to_image.py | 1 - .../pipelines/latent_diffusion/pipeline_latent_diffusion.py | 1 - .../pipeline_latent_diffusion_superresolution.py | 1 - src/diffusers/pipelines/stable_audio/modeling_stable_audio.py | 1 - 17 files changed, 17 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py index 7b0f9889a5..dc5e775f67 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py @@ -18,7 +18,6 @@ import numpy as np import torch import torch.nn as nn import torch.nn.functional as F -import torch.utils.checkpoint from ...configuration_utils import ConfigMixin, register_to_config from ...utils import logging diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py index 87ac406592..9872cf0968 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py @@ -23,7 +23,6 @@ from typing import List, Optional, Tuple, Union import torch import torch.nn as nn import torch.nn.functional as F -import torch.utils.checkpoint from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_wan.py b/src/diffusers/models/autoencoders/autoencoder_kl_wan.py index e6e58c1cce..f95c4cf374 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_wan.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_wan.py @@ -17,7 +17,6 @@ from typing import List, Optional, Tuple, Union import torch import torch.nn as nn import torch.nn.functional as F -import torch.utils.checkpoint from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin diff --git a/src/diffusers/models/controlnets/controlnet_xs.py b/src/diffusers/models/controlnets/controlnet_xs.py index bcb4e25986..f5c69b9a46 100644 --- a/src/diffusers/models/controlnets/controlnet_xs.py +++ b/src/diffusers/models/controlnets/controlnet_xs.py @@ -16,7 +16,6 @@ from math import gcd from typing import Any, Dict, List, Optional, Tuple, Union import torch -import torch.utils.checkpoint from torch import Tensor, nn from ...configuration_utils import ConfigMixin, register_to_config diff --git a/src/diffusers/models/transformers/stable_audio_transformer.py b/src/diffusers/models/transformers/stable_audio_transformer.py index 969e6db122..ac9b3fca41 100644 --- a/src/diffusers/models/transformers/stable_audio_transformer.py +++ b/src/diffusers/models/transformers/stable_audio_transformer.py @@ -18,7 +18,6 @@ from typing import Dict, Optional, Union import numpy as np import torch import torch.nn as nn -import torch.utils.checkpoint from ...configuration_utils import ConfigMixin, register_to_config from ...utils import logging diff --git a/src/diffusers/models/unets/unet_2d_condition.py b/src/diffusers/models/unets/unet_2d_condition.py index 33bda8cb1e..f04d3dfa01 100644 --- a/src/diffusers/models/unets/unet_2d_condition.py +++ b/src/diffusers/models/unets/unet_2d_condition.py @@ -16,7 +16,6 @@ from typing import Any, Dict, List, Optional, Tuple, Union import torch import torch.nn as nn -import torch.utils.checkpoint from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import PeftAdapterMixin, UNet2DConditionLoadersMixin diff --git a/src/diffusers/models/unets/unet_3d_condition.py b/src/diffusers/models/unets/unet_3d_condition.py index b5151f3c9a..6a119185b8 100644 --- a/src/diffusers/models/unets/unet_3d_condition.py +++ b/src/diffusers/models/unets/unet_3d_condition.py @@ -18,7 +18,6 @@ from typing import Any, Dict, List, Optional, Tuple, Union import torch import torch.nn as nn -import torch.utils.checkpoint from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import UNet2DConditionLoadersMixin diff --git a/src/diffusers/models/unets/unet_i2vgen_xl.py b/src/diffusers/models/unets/unet_i2vgen_xl.py index 7148723a84..3dba8edca7 100644 --- a/src/diffusers/models/unets/unet_i2vgen_xl.py +++ b/src/diffusers/models/unets/unet_i2vgen_xl.py @@ -16,7 +16,6 @@ from typing import Any, Dict, Optional, Tuple, Union import torch import torch.nn as nn -import torch.utils.checkpoint from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import UNet2DConditionLoadersMixin diff --git a/src/diffusers/models/unets/unet_kandinsky3.py b/src/diffusers/models/unets/unet_kandinsky3.py index 423669a22f..27241ce2e6 100644 --- a/src/diffusers/models/unets/unet_kandinsky3.py +++ b/src/diffusers/models/unets/unet_kandinsky3.py @@ -16,7 +16,6 @@ from dataclasses import dataclass from typing import Dict, Tuple, Union import torch -import torch.utils.checkpoint from torch import nn from ...configuration_utils import ConfigMixin, register_to_config diff --git a/src/diffusers/models/unets/unet_motion_model.py b/src/diffusers/models/unets/unet_motion_model.py index 26616e53bd..18d5eb917f 100644 --- a/src/diffusers/models/unets/unet_motion_model.py +++ b/src/diffusers/models/unets/unet_motion_model.py @@ -18,7 +18,6 @@ from typing import Any, Dict, Optional, Tuple, Union import torch import torch.nn as nn import torch.nn.functional as F -import torch.utils.checkpoint from ...configuration_utils import ConfigMixin, FrozenDict, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin, UNet2DConditionLoadersMixin diff --git a/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py b/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py index 546ae9239a..b6b40cd6e6 100644 --- a/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py +++ b/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py @@ -17,7 +17,6 @@ from typing import Any, Dict, List, Optional, Tuple, Union import torch import torch.nn as nn -import torch.utils.checkpoint from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import UNet2DConditionLoadersMixin diff --git a/src/diffusers/pipelines/blip_diffusion/modeling_blip2.py b/src/diffusers/pipelines/blip_diffusion/modeling_blip2.py index 928698e442..b061ac2636 100644 --- a/src/diffusers/pipelines/blip_diffusion/modeling_blip2.py +++ b/src/diffusers/pipelines/blip_diffusion/modeling_blip2.py @@ -14,7 +14,6 @@ from typing import Optional, Tuple, Union import torch -import torch.utils.checkpoint from torch import nn from transformers import BertTokenizer from transformers.activations import QuickGELUActivation as QuickGELU diff --git a/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_image_variation.py b/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_image_variation.py index 2beb0be57b..034a022641 100644 --- a/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_image_variation.py +++ b/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_image_variation.py @@ -18,7 +18,6 @@ from typing import Callable, List, Optional, Union import numpy as np import PIL.Image import torch -import torch.utils.checkpoint from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection from ....image_processor import VaeImageProcessor diff --git a/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_text_to_image.py b/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_text_to_image.py index adfd899e76..2f54f4fc98 100644 --- a/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_text_to_image.py +++ b/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_text_to_image.py @@ -16,7 +16,6 @@ import inspect from typing import Callable, List, Optional, Union import torch -import torch.utils.checkpoint from transformers import CLIPImageProcessor, CLIPTextModelWithProjection, CLIPTokenizer from ....image_processor import VaeImageProcessor diff --git a/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py b/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py index bc50835d19..f1bf4701e3 100644 --- a/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py +++ b/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py @@ -17,7 +17,6 @@ from typing import List, Optional, Tuple, Union import torch import torch.nn as nn -import torch.utils.checkpoint from transformers import PretrainedConfig, PreTrainedModel, PreTrainedTokenizer from transformers.activations import ACT2FN from transformers.modeling_outputs import BaseModelOutput diff --git a/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion_superresolution.py b/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion_superresolution.py index 273e97f1ec..631539e5c6 100644 --- a/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion_superresolution.py +++ b/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion_superresolution.py @@ -4,7 +4,6 @@ from typing import List, Optional, Tuple, Union import numpy as np import PIL.Image import torch -import torch.utils.checkpoint from ...models import UNet2DModel, VQModel from ...schedulers import ( diff --git a/src/diffusers/pipelines/stable_audio/modeling_stable_audio.py b/src/diffusers/pipelines/stable_audio/modeling_stable_audio.py index 89d4d2dca5..07b382dfc4 100644 --- a/src/diffusers/pipelines/stable_audio/modeling_stable_audio.py +++ b/src/diffusers/pipelines/stable_audio/modeling_stable_audio.py @@ -18,7 +18,6 @@ from typing import Optional import torch import torch.nn as nn -import torch.utils.checkpoint from ...configuration_utils import ConfigMixin, register_to_config from ...models.modeling_utils import ModelMixin From aed636f5f089754425614872e0b4781dab65179f Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Wed, 15 Oct 2025 10:57:09 +0530 Subject: [PATCH 067/123] [tests] fix clapconfig for text backbone in audioldm2 (#12490) fix clapconfig for text backbone in audioldm2 --- tests/pipelines/audioldm2/test_audioldm2.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/pipelines/audioldm2/test_audioldm2.py b/tests/pipelines/audioldm2/test_audioldm2.py index e4bc5cc110..14ff1272a2 100644 --- a/tests/pipelines/audioldm2/test_audioldm2.py +++ b/tests/pipelines/audioldm2/test_audioldm2.py @@ -138,10 +138,8 @@ class AudioLDM2PipelineFastTests(PipelineTesterMixin, unittest.TestCase): patch_stride=2, patch_embed_input_channels=4, ) - text_encoder_config = ClapConfig.from_text_audio_configs( - text_config=text_branch_config, - audio_config=audio_branch_config, - projection_dim=16, + text_encoder_config = ClapConfig( + text_config=text_branch_config, audio_config=audio_branch_config, projection_dim=16 ) text_encoder = ClapModel(text_encoder_config) tokenizer = RobertaTokenizer.from_pretrained("hf-internal-testing/tiny-random-roberta", model_max_length=77) From f072c64bf22c6a259d38435b19a8ef3eb67c96a8 Mon Sep 17 00:00:00 2001 From: YiYi Xu Date: Wed, 15 Oct 2025 07:41:17 -1000 Subject: [PATCH 068/123] ltx0.9.8 (without IC lora, autoregressive sampling) (#12493) update Co-authored-by: Aryan --- docs/source/en/api/pipelines/ltx_video.md | 93 ++++++++++++++++++- scripts/convert_ltx_to_diffusers.py | 11 ++- .../ltx/pipeline_ltx_latent_upsample.py | 42 ++++++++- 3 files changed, 142 insertions(+), 4 deletions(-) diff --git a/docs/source/en/api/pipelines/ltx_video.md b/docs/source/en/api/pipelines/ltx_video.md index d87c57ced7..940144538a 100644 --- a/docs/source/en/api/pipelines/ltx_video.md +++ b/docs/source/en/api/pipelines/ltx_video.md @@ -254,8 +254,8 @@ export_to_video(video, "output.mp4", fps=24) pipeline.vae.enable_tiling() def round_to_nearest_resolution_acceptable_by_vae(height, width): - height = height - (height % pipeline.vae_temporal_compression_ratio) - width = width - (width % pipeline.vae_temporal_compression_ratio) + height = height - (height % pipeline.vae_spatial_compression_ratio) + width = width - (width % pipeline.vae_spatial_compression_ratio) return height, width prompt = """ @@ -325,6 +325,95 @@ export_to_video(video, "output.mp4", fps=24) +- LTX-Video 0.9.8 distilled model is similar to the 0.9.7 variant. It is guidance and timestep-distilled, and similar inference code can be used as above. An improvement of this version is that it supports generating very long videos. Additionally, it supports using tone mapping to improve the quality of the generated video using the `tone_map_compression_ratio` parameter. The default value of `0.6` is recommended. + +
+ Show example code + + ```python + import torch + from diffusers import LTXConditionPipeline, LTXLatentUpsamplePipeline + from diffusers.pipelines.ltx.pipeline_ltx_condition import LTXVideoCondition + from diffusers.pipelines.ltx.modeling_latent_upsampler import LTXLatentUpsamplerModel + from diffusers.utils import export_to_video, load_video + + pipeline = LTXConditionPipeline.from_pretrained("Lightricks/LTX-Video-0.9.8-13B-distilled", torch_dtype=torch.bfloat16) + # TODO: Update the checkpoint here once updated in LTX org + upsampler = LTXLatentUpsamplerModel.from_pretrained("a-r-r-o-w/LTX-0.9.8-Latent-Upsampler", torch_dtype=torch.bfloat16) + pipe_upsample = LTXLatentUpsamplePipeline(vae=pipeline.vae, latent_upsampler=upsampler).to(torch.bfloat16) + pipeline.to("cuda") + pipe_upsample.to("cuda") + pipeline.vae.enable_tiling() + + def round_to_nearest_resolution_acceptable_by_vae(height, width): + height = height - (height % pipeline.vae_spatial_compression_ratio) + width = width - (width % pipeline.vae_spatial_compression_ratio) + return height, width + + prompt = """The camera pans over a snow-covered mountain range, revealing a vast expanse of snow-capped peaks and valleys.The mountains are covered in a thick layer of snow, with some areas appearing almost white while others have a slightly darker, almost grayish hue. The peaks are jagged and irregular, with some rising sharply into the sky while others are more rounded. The valleys are deep and narrow, with steep slopes that are also covered in snow. The trees in the foreground are mostly bare, with only a few leaves remaining on their branches. The sky is overcast, with thick clouds obscuring the sun. The overall impression is one of peace and tranquility, with the snow-covered mountains standing as a testament to the power and beauty of nature.""" + # prompt = """A woman walks away from a white Jeep parked on a city street at night, then ascends a staircase and knocks on a door. The woman, wearing a dark jacket and jeans, walks away from the Jeep parked on the left side of the street, her back to the camera; she walks at a steady pace, her arms swinging slightly by her sides; the street is dimly lit, with streetlights casting pools of light on the wet pavement; a man in a dark jacket and jeans walks past the Jeep in the opposite direction; the camera follows the woman from behind as she walks up a set of stairs towards a building with a green door; she reaches the top of the stairs and turns left, continuing to walk towards the building; she reaches the door and knocks on it with her right hand; the camera remains stationary, focused on the doorway; the scene is captured in real-life footage.""" + negative_prompt = "bright colors, symbols, graffiti, watermarks, worst quality, inconsistent motion, blurry, jittery, distorted" + expected_height, expected_width = 480, 832 + downscale_factor = 2 / 3 + # num_frames = 161 + num_frames = 361 + + # 1. Generate video at smaller resolution + downscaled_height, downscaled_width = int(expected_height * downscale_factor), int(expected_width * downscale_factor) + downscaled_height, downscaled_width = round_to_nearest_resolution_acceptable_by_vae(downscaled_height, downscaled_width) + latents = pipeline( + prompt=prompt, + negative_prompt=negative_prompt, + width=downscaled_width, + height=downscaled_height, + num_frames=num_frames, + timesteps=[1000, 993, 987, 981, 975, 909, 725, 0.03], + decode_timestep=0.05, + decode_noise_scale=0.025, + image_cond_noise_scale=0.0, + guidance_scale=1.0, + guidance_rescale=0.7, + generator=torch.Generator().manual_seed(0), + output_type="latent", + ).frames + + # 2. Upscale generated video using latent upsampler with fewer inference steps + # The available latent upsampler upscales the height/width by 2x + upscaled_height, upscaled_width = downscaled_height * 2, downscaled_width * 2 + upscaled_latents = pipe_upsample( + latents=latents, + adain_factor=1.0, + tone_map_compression_ratio=0.6, + output_type="latent" + ).frames + + # 3. Denoise the upscaled video with few steps to improve texture (optional, but recommended) + video = pipeline( + prompt=prompt, + negative_prompt=negative_prompt, + width=upscaled_width, + height=upscaled_height, + num_frames=num_frames, + denoise_strength=0.999, # Effectively, 4 inference steps out of 5 + timesteps=[1000, 909, 725, 421, 0], + latents=upscaled_latents, + decode_timestep=0.05, + decode_noise_scale=0.025, + image_cond_noise_scale=0.0, + guidance_scale=1.0, + guidance_rescale=0.7, + generator=torch.Generator().manual_seed(0), + output_type="pil", + ).frames[0] + + # 4. Downscale the video to the expected resolution + video = [frame.resize((expected_width, expected_height)) for frame in video] + + export_to_video(video, "output.mp4", fps=24) + ``` + +
+ - LTX-Video supports LoRAs with [`~loaders.LTXVideoLoraLoaderMixin.load_lora_weights`].
diff --git a/scripts/convert_ltx_to_diffusers.py b/scripts/convert_ltx_to_diffusers.py index 256312cc72..19e5602039 100644 --- a/scripts/convert_ltx_to_diffusers.py +++ b/scripts/convert_ltx_to_diffusers.py @@ -369,6 +369,15 @@ def get_spatial_latent_upsampler_config(version: str) -> Dict[str, Any]: "spatial_upsample": True, "temporal_upsample": False, } + elif version == "0.9.8": + config = { + "in_channels": 128, + "mid_channels": 512, + "num_blocks_per_stage": 4, + "dims": 3, + "spatial_upsample": True, + "temporal_upsample": False, + } else: raise ValueError(f"Unsupported version: {version}") return config @@ -402,7 +411,7 @@ def get_args(): "--version", type=str, default="0.9.0", - choices=["0.9.0", "0.9.1", "0.9.5", "0.9.7"], + choices=["0.9.0", "0.9.1", "0.9.5", "0.9.7", "0.9.8"], help="Version of the LTX model", ) return parser.parse_args() diff --git a/src/diffusers/pipelines/ltx/pipeline_ltx_latent_upsample.py b/src/diffusers/pipelines/ltx/pipeline_ltx_latent_upsample.py index 1e94f6895f..9acff105e5 100644 --- a/src/diffusers/pipelines/ltx/pipeline_ltx_latent_upsample.py +++ b/src/diffusers/pipelines/ltx/pipeline_ltx_latent_upsample.py @@ -121,6 +121,38 @@ class LTXLatentUpsamplePipeline(DiffusionPipeline): result = torch.lerp(latents, result, factor) return result + def tone_map_latents(self, latents: torch.Tensor, compression: float) -> torch.Tensor: + """ + Applies a non-linear tone-mapping function to latent values to reduce their dynamic range in a perceptually + smooth way using a sigmoid-based compression. + + This is useful for regularizing high-variance latents or for conditioning outputs during generation, especially + when controlling dynamic behavior with a `compression` factor. + + Args: + latents : torch.Tensor + Input latent tensor with arbitrary shape. Expected to be roughly in [-1, 1] or [0, 1] range. + compression : float + Compression strength in the range [0, 1]. + - 0.0: No tone-mapping (identity transform) + - 1.0: Full compression effect + + Returns: + torch.Tensor + The tone-mapped latent tensor of the same shape as input. + """ + # Remap [0-1] to [0-0.75] and apply sigmoid compression in one shot + scale_factor = compression * 0.75 + abs_latents = torch.abs(latents) + + # Sigmoid compression: sigmoid shifts large values toward 0.2, small values stay ~1.0 + # When scale_factor=0, sigmoid term vanishes, when scale_factor=0.75, full effect + sigmoid_term = torch.sigmoid(4.0 * scale_factor * (abs_latents - 1.0)) + scales = 1.0 - 0.8 * scale_factor * sigmoid_term + + filtered = latents * scales + return filtered + @staticmethod # Copied from diffusers.pipelines.ltx.pipeline_ltx.LTXPipeline._normalize_latents def _normalize_latents( @@ -196,7 +228,7 @@ class LTXLatentUpsamplePipeline(DiffusionPipeline): ) self.vae.disable_tiling() - def check_inputs(self, video, height, width, latents): + def check_inputs(self, video, height, width, latents, tone_map_compression_ratio): if height % self.vae_spatial_compression_ratio != 0 or width % self.vae_spatial_compression_ratio != 0: raise ValueError(f"`height` and `width` have to be divisible by 32 but are {height} and {width}.") @@ -205,6 +237,9 @@ class LTXLatentUpsamplePipeline(DiffusionPipeline): if video is None and latents is None: raise ValueError("One of `video` or `latents` has to be provided.") + if not (0 <= tone_map_compression_ratio <= 1): + raise ValueError("`tone_map_compression_ratio` must be in the range [0, 1]") + @torch.no_grad() def __call__( self, @@ -215,6 +250,7 @@ class LTXLatentUpsamplePipeline(DiffusionPipeline): decode_timestep: Union[float, List[float]] = 0.0, decode_noise_scale: Optional[Union[float, List[float]]] = None, adain_factor: float = 0.0, + tone_map_compression_ratio: float = 0.0, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, output_type: Optional[str] = "pil", return_dict: bool = True, @@ -224,6 +260,7 @@ class LTXLatentUpsamplePipeline(DiffusionPipeline): height=height, width=width, latents=latents, + tone_map_compression_ratio=tone_map_compression_ratio, ) if video is not None: @@ -266,6 +303,9 @@ class LTXLatentUpsamplePipeline(DiffusionPipeline): else: latents = latents_upsampled + if tone_map_compression_ratio > 0.0: + latents = self.tone_map_latents(latents, tone_map_compression_ratio) + if output_type == "latent": latents = self._normalize_latents( latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor From 26475082cb9db5da659378c978496999d338b960 Mon Sep 17 00:00:00 2001 From: Steven Liu <59462357+stevhliu@users.noreply.github.com> Date: Thu, 16 Oct 2025 09:19:30 -0700 Subject: [PATCH 069/123] [docs] Attention checks (#12486) * checks * feedback --------- Co-authored-by: Sayak Paul --- .../en/optimization/attention_backends.md | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/docs/source/en/optimization/attention_backends.md b/docs/source/en/optimization/attention_backends.md index e603878a63..8be2c06030 100644 --- a/docs/source/en/optimization/attention_backends.md +++ b/docs/source/en/optimization/attention_backends.md @@ -81,6 +81,45 @@ with attention_backend("_flash_3_hub"): > [!TIP] > Most attention backends support `torch.compile` without graph breaks and can be used to further speed up inference. +## Checks + +The attention dispatcher includes debugging checks that catch common errors before they cause problems. + +1. Device checks verify that query, key, and value tensors live on the same device. +2. Data type checks confirm tensors have matching dtypes and use either bfloat16 or float16. +3. Shape checks validate tensor dimensions and prevent mixing attention masks with causal flags. + +Enable these checks by setting the `DIFFUSERS_ATTN_CHECKS` environment variable. Checks add overhead to every attention operation, so they're disabled by default. + +```bash +export DIFFUSERS_ATTN_CHECKS=yes +``` + +The checks are run now before every attention operation. + +```py +import torch + +query = torch.randn(1, 10, 8, 64, dtype=torch.bfloat16, device="cuda") +key = torch.randn(1, 10, 8, 64, dtype=torch.bfloat16, device="cuda") +value = torch.randn(1, 10, 8, 64, dtype=torch.bfloat16, device="cuda") + +try: + with attention_backend("flash"): + output = dispatch_attention_fn(query, key, value) + print("✓ Flash Attention works with checks enabled") +except Exception as e: + print(f"✗ Flash Attention failed: {e}") +``` + +You can also configure the registry directly. + +```py +from diffusers.models.attention_dispatch import _AttentionBackendRegistry + +_AttentionBackendRegistry._checks_enabled = True +``` + ## Available backends Refer to the table below for a complete list of available attention backends and their variants. From dbe413668dbf9d527944b51ef5f99a60b36a90be Mon Sep 17 00:00:00 2001 From: Steven Liu <59462357+stevhliu@users.noreply.github.com> Date: Thu, 16 Oct 2025 10:38:16 -0700 Subject: [PATCH 070/123] [CI] Check links (#12491) * check links * update * feedback * remove --- .github/workflows/build_pr_documentation.yml | 26 ++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/.github/workflows/build_pr_documentation.yml b/.github/workflows/build_pr_documentation.yml index 52e0757331..f47645c1f6 100644 --- a/.github/workflows/build_pr_documentation.yml +++ b/.github/workflows/build_pr_documentation.yml @@ -12,7 +12,33 @@ concurrency: cancel-in-progress: true jobs: + check-links: + runs-on: ubuntu-latest + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.10' + + - name: Install uv + run: | + curl -LsSf https://astral.sh/uv/install.sh | sh + echo "$HOME/.cargo/bin" >> $GITHUB_PATH + + - name: Install doc-builder + run: | + uv pip install --system git+https://github.com/huggingface/doc-builder.git@main + + - name: Check documentation links + run: | + uv run doc-builder check-links docs/source/en + build: + needs: check-links uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@main with: commit_sha: ${{ github.event.pull_request.head.sha }} From 4715c5c769e5ed2f3f174ee7aa76c67eb875b669 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Fri, 17 Oct 2025 10:35:19 +0530 Subject: [PATCH 071/123] [ci] xfail more incorrect transformer imports. (#12455) * xfail more incorrect transformer imports. * xfail more. * up * up * up --- tests/pipelines/test_pipelines.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/pipelines/test_pipelines.py b/tests/pipelines/test_pipelines.py index 5b86423553..a17db3ff0c 100644 --- a/tests/pipelines/test_pipelines.py +++ b/tests/pipelines/test_pipelines.py @@ -582,7 +582,7 @@ class DownloadTests(unittest.TestCase): assert not any(f.endswith(unexpected_ext) for f in files) assert all(variant in f for f in model_files if f.endswith(model_ext) and variant is not None) - @pytest.mark.xfail(condition=is_transformers_version(">", "4.56.2"), reason="Some import error", strict=True) + @pytest.mark.xfail(condition=is_transformers_version(">", "4.56.2"), reason="Some import error", strict=False) def test_download_legacy_variants_with_sharded_ckpts_raises_warning(self): repo_id = "hf-internal-testing/tiny-stable-diffusion-pipe-variants-all-kinds" logger = logging.get_logger("diffusers.pipelines.pipeline_utils") @@ -629,6 +629,7 @@ class DownloadTests(unittest.TestCase): # https://huggingface.co/hf-internal-testing/stable-diffusion-broken-variants/tree/main/unet assert len(files) == 15, f"We should only download 15 files, not {len(files)}" + @pytest.mark.xfail(condition=is_transformers_version(">", "4.56.2"), reason="Some import error", strict=False) def test_download_bin_only_variant_exists_for_model(self): variant = None use_safetensors = False @@ -674,6 +675,7 @@ class DownloadTests(unittest.TestCase): assert "Could not find the necessary `safetensors` weights" in str(error_context.exception) + @pytest.mark.xfail(condition=is_transformers_version(">", "4.56.2"), reason="Some import error", strict=False) def test_download_bin_variant_does_not_exist_for_model(self): variant = "no_ema" use_safetensors = False @@ -689,7 +691,7 @@ class DownloadTests(unittest.TestCase): ) assert "Error no file name" in str(error_context.exception) - @pytest.mark.xfail(condition=is_transformers_version(">", "4.56.2"), reason="Some import error", strict=True) + @pytest.mark.xfail(condition=is_transformers_version(">", "4.56.2"), reason="Some import error", strict=False) def test_local_save_load_index(self): prompt = "hello" for variant in [None, "fp16"]: @@ -1584,7 +1586,7 @@ class PipelineFastTests(unittest.TestCase): assert pipeline.scheduler is not None assert pipeline.feature_extractor is not None - @pytest.mark.xfail(condition=is_transformers_version(">", "4.56.2"), reason="Some import error", strict=True) + @pytest.mark.xfail(condition=is_transformers_version(">", "4.56.2"), reason="Some import error", strict=False) def test_no_pytorch_download_when_doing_safetensors(self): # by default we don't download with tempfile.TemporaryDirectory() as tmpdirname: @@ -1604,7 +1606,7 @@ class PipelineFastTests(unittest.TestCase): # pytorch does not assert not os.path.exists(os.path.join(path, "diffusion_pytorch_model.bin")) - @pytest.mark.xfail(condition=is_transformers_version(">", "4.56.2"), reason="Some import error", strict=True) + @pytest.mark.xfail(condition=is_transformers_version(">", "4.56.2"), reason="Some import error", strict=False) def test_no_safetensors_download_when_doing_pytorch(self): use_safetensors = False @@ -1890,6 +1892,7 @@ class PipelineFastTests(unittest.TestCase): "DDUF/tiny-flux-dev-pipe-dduf", dduf_file="fluxpipeline.dduf", load_connected_pipeline=True ) + @pytest.mark.xfail(condition=is_transformers_version(">", "4.56.2"), reason="Some import error", strict=False) def test_wrong_model(self): tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") with self.assertRaises(ValueError) as error_context: From af769881d37fe916afef2c47279f66c79f5f2714 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Fri, 17 Oct 2025 12:02:29 +0530 Subject: [PATCH 072/123] [tests] introduce `VAETesterMixin` to consolidate tests for slicing and tiling (#12374) * up * up * up * up * up * u[ * up * up * up --- .../test_models_asymmetric_autoencoder_kl.py | 5 +- .../test_models_autoencoder_cosmos.py | 9 +- .../test_models_autoencoder_dc.py | 9 +- .../test_models_autoencoder_hunyuan_video.py | 73 +-------- .../test_models_autoencoder_kl.py | 67 +-------- .../test_models_autoencoder_kl_cogvideox.py | 67 +-------- ..._models_autoencoder_kl_temporal_decoder.py | 9 +- .../test_models_autoencoder_ltx_video.py | 38 +---- .../test_models_autoencoder_magvit.py | 11 +- .../test_models_autoencoder_mochi.py | 19 +-- .../test_models_autoencoder_oobleck.py | 9 +- .../test_models_autoencoder_tiny.py | 36 +---- .../test_models_autoencoder_wan.py | 69 +-------- .../test_models_consistency_decoder_vae.py | 67 +-------- tests/models/autoencoders/test_models_vq.py | 12 +- tests/models/autoencoders/testing_utils.py | 142 ++++++++++++++++++ tests/models/test_modeling_common.py | 8 + 17 files changed, 204 insertions(+), 446 deletions(-) create mode 100644 tests/models/autoencoders/testing_utils.py diff --git a/tests/models/autoencoders/test_models_asymmetric_autoencoder_kl.py b/tests/models/autoencoders/test_models_asymmetric_autoencoder_kl.py index 7eb830cd50..2476ab92f7 100644 --- a/tests/models/autoencoders/test_models_asymmetric_autoencoder_kl.py +++ b/tests/models/autoencoders/test_models_asymmetric_autoencoder_kl.py @@ -35,13 +35,14 @@ from ...testing_utils import ( torch_all_close, torch_device, ) -from ..test_modeling_common import ModelTesterMixin, UNetTesterMixin +from ..test_modeling_common import ModelTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() -class AutoencoderKLTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): +class AsymmetricAutoencoderKLTests(ModelTesterMixin, AutoencoderTesterMixin, unittest.TestCase): model_class = AsymmetricAutoencoderKL main_input_name = "sample" base_precision = 1e-2 diff --git a/tests/models/autoencoders/test_models_autoencoder_cosmos.py b/tests/models/autoencoders/test_models_autoencoder_cosmos.py index ceccc2364e..5898ae776a 100644 --- a/tests/models/autoencoders/test_models_autoencoder_cosmos.py +++ b/tests/models/autoencoders/test_models_autoencoder_cosmos.py @@ -17,13 +17,14 @@ import unittest from diffusers import AutoencoderKLCosmos from ...testing_utils import enable_full_determinism, floats_tensor, torch_device -from ..test_modeling_common import ModelTesterMixin, UNetTesterMixin +from ..test_modeling_common import ModelTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() -class AutoencoderKLCosmosTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): +class AutoencoderKLCosmosTests(ModelTesterMixin, AutoencoderTesterMixin, unittest.TestCase): model_class = AutoencoderKLCosmos main_input_name = "sample" base_precision = 1e-2 @@ -80,7 +81,3 @@ class AutoencoderKLCosmosTests(ModelTesterMixin, UNetTesterMixin, unittest.TestC @unittest.skip("Not sure why this test fails. Investigate later.") def test_effective_gradient_checkpointing(self): pass - - @unittest.skip("Unsupported test.") - def test_forward_with_norm_groups(self): - pass diff --git a/tests/models/autoencoders/test_models_autoencoder_dc.py b/tests/models/autoencoders/test_models_autoencoder_dc.py index 56f172f1c8..a6912f3eba 100644 --- a/tests/models/autoencoders/test_models_autoencoder_dc.py +++ b/tests/models/autoencoders/test_models_autoencoder_dc.py @@ -22,13 +22,14 @@ from ...testing_utils import ( floats_tensor, torch_device, ) -from ..test_modeling_common import ModelTesterMixin, UNetTesterMixin +from ..test_modeling_common import ModelTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() -class AutoencoderDCTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): +class AutoencoderDCTests(ModelTesterMixin, AutoencoderTesterMixin, unittest.TestCase): model_class = AutoencoderDC main_input_name = "sample" base_precision = 1e-2 @@ -81,7 +82,3 @@ class AutoencoderDCTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): init_dict = self.get_autoencoder_dc_config() inputs_dict = self.dummy_input return init_dict, inputs_dict - - @unittest.skip("AutoencoderDC does not support `norm_num_groups` because it does not use GroupNorm.") - def test_forward_with_norm_groups(self): - pass diff --git a/tests/models/autoencoders/test_models_autoencoder_hunyuan_video.py b/tests/models/autoencoders/test_models_autoencoder_hunyuan_video.py index 6f91f8bfa9..9813772a7c 100644 --- a/tests/models/autoencoders/test_models_autoencoder_hunyuan_video.py +++ b/tests/models/autoencoders/test_models_autoencoder_hunyuan_video.py @@ -20,18 +20,15 @@ import torch from diffusers import AutoencoderKLHunyuanVideo from diffusers.models.autoencoders.autoencoder_kl_hunyuan_video import prepare_causal_attention_mask -from ...testing_utils import ( - enable_full_determinism, - floats_tensor, - torch_device, -) -from ..test_modeling_common import ModelTesterMixin, UNetTesterMixin +from ...testing_utils import enable_full_determinism, floats_tensor, torch_device +from ..test_modeling_common import ModelTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() -class AutoencoderKLHunyuanVideoTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): +class AutoencoderKLHunyuanVideoTests(ModelTesterMixin, AutoencoderTesterMixin, unittest.TestCase): model_class = AutoencoderKLHunyuanVideo main_input_name = "sample" base_precision = 1e-2 @@ -87,68 +84,6 @@ class AutoencoderKLHunyuanVideoTests(ModelTesterMixin, UNetTesterMixin, unittest inputs_dict = self.dummy_input return init_dict, inputs_dict - def test_enable_disable_tiling(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - - torch.manual_seed(0) - model = self.model_class(**init_dict).to(torch_device) - - inputs_dict.update({"return_dict": False}) - - torch.manual_seed(0) - output_without_tiling = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - torch.manual_seed(0) - model.enable_tiling() - output_with_tiling = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - self.assertLess( - (output_without_tiling.detach().cpu().numpy() - output_with_tiling.detach().cpu().numpy()).max(), - 0.5, - "VAE tiling should not affect the inference results", - ) - - torch.manual_seed(0) - model.disable_tiling() - output_without_tiling_2 = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - self.assertEqual( - output_without_tiling.detach().cpu().numpy().all(), - output_without_tiling_2.detach().cpu().numpy().all(), - "Without tiling outputs should match with the outputs when tiling is manually disabled.", - ) - - def test_enable_disable_slicing(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - - torch.manual_seed(0) - model = self.model_class(**init_dict).to(torch_device) - - inputs_dict.update({"return_dict": False}) - - torch.manual_seed(0) - output_without_slicing = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - torch.manual_seed(0) - model.enable_slicing() - output_with_slicing = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - self.assertLess( - (output_without_slicing.detach().cpu().numpy() - output_with_slicing.detach().cpu().numpy()).max(), - 0.5, - "VAE slicing should not affect the inference results", - ) - - torch.manual_seed(0) - model.disable_slicing() - output_without_slicing_2 = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - self.assertEqual( - output_without_slicing.detach().cpu().numpy().all(), - output_without_slicing_2.detach().cpu().numpy().all(), - "Without slicing outputs should match with the outputs when slicing is manually disabled.", - ) - def test_gradient_checkpointing_is_applied(self): expected_set = { "HunyuanVideoDecoder3D", diff --git a/tests/models/autoencoders/test_models_autoencoder_kl.py b/tests/models/autoencoders/test_models_autoencoder_kl.py index 662a3f1b80..5f11c6cb0a 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl.py @@ -35,13 +35,14 @@ from ...testing_utils import ( torch_all_close, torch_device, ) -from ..test_modeling_common import ModelTesterMixin, UNetTesterMixin +from ..test_modeling_common import ModelTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() -class AutoencoderKLTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): +class AutoencoderKLTests(ModelTesterMixin, AutoencoderTesterMixin, unittest.TestCase): model_class = AutoencoderKL main_input_name = "sample" base_precision = 1e-2 @@ -83,68 +84,6 @@ class AutoencoderKLTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): inputs_dict = self.dummy_input return init_dict, inputs_dict - def test_enable_disable_tiling(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - - torch.manual_seed(0) - model = self.model_class(**init_dict).to(torch_device) - - inputs_dict.update({"return_dict": False}) - - torch.manual_seed(0) - output_without_tiling = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - torch.manual_seed(0) - model.enable_tiling() - output_with_tiling = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - self.assertLess( - (output_without_tiling.detach().cpu().numpy() - output_with_tiling.detach().cpu().numpy()).max(), - 0.5, - "VAE tiling should not affect the inference results", - ) - - torch.manual_seed(0) - model.disable_tiling() - output_without_tiling_2 = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - self.assertEqual( - output_without_tiling.detach().cpu().numpy().all(), - output_without_tiling_2.detach().cpu().numpy().all(), - "Without tiling outputs should match with the outputs when tiling is manually disabled.", - ) - - def test_enable_disable_slicing(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - - torch.manual_seed(0) - model = self.model_class(**init_dict).to(torch_device) - - inputs_dict.update({"return_dict": False}) - - torch.manual_seed(0) - output_without_slicing = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - torch.manual_seed(0) - model.enable_slicing() - output_with_slicing = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - self.assertLess( - (output_without_slicing.detach().cpu().numpy() - output_with_slicing.detach().cpu().numpy()).max(), - 0.5, - "VAE slicing should not affect the inference results", - ) - - torch.manual_seed(0) - model.disable_slicing() - output_without_slicing_2 = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - self.assertEqual( - output_without_slicing.detach().cpu().numpy().all(), - output_without_slicing_2.detach().cpu().numpy().all(), - "Without slicing outputs should match with the outputs when slicing is manually disabled.", - ) - def test_gradient_checkpointing_is_applied(self): expected_set = {"Decoder", "Encoder", "UNetMidBlock2D"} super().test_gradient_checkpointing_is_applied(expected_set=expected_set) diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_cogvideox.py b/tests/models/autoencoders/test_models_autoencoder_kl_cogvideox.py index 739daf2a49..b6d59489d9 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_cogvideox.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_cogvideox.py @@ -24,13 +24,14 @@ from ...testing_utils import ( floats_tensor, torch_device, ) -from ..test_modeling_common import ModelTesterMixin, UNetTesterMixin +from ..test_modeling_common import ModelTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() -class AutoencoderKLCogVideoXTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): +class AutoencoderKLCogVideoXTests(ModelTesterMixin, AutoencoderTesterMixin, unittest.TestCase): model_class = AutoencoderKLCogVideoX main_input_name = "sample" base_precision = 1e-2 @@ -82,68 +83,6 @@ class AutoencoderKLCogVideoXTests(ModelTesterMixin, UNetTesterMixin, unittest.Te inputs_dict = self.dummy_input return init_dict, inputs_dict - def test_enable_disable_tiling(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - - torch.manual_seed(0) - model = self.model_class(**init_dict).to(torch_device) - - inputs_dict.update({"return_dict": False}) - - torch.manual_seed(0) - output_without_tiling = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - torch.manual_seed(0) - model.enable_tiling() - output_with_tiling = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - self.assertLess( - (output_without_tiling.detach().cpu().numpy() - output_with_tiling.detach().cpu().numpy()).max(), - 0.5, - "VAE tiling should not affect the inference results", - ) - - torch.manual_seed(0) - model.disable_tiling() - output_without_tiling_2 = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - self.assertEqual( - output_without_tiling.detach().cpu().numpy().all(), - output_without_tiling_2.detach().cpu().numpy().all(), - "Without tiling outputs should match with the outputs when tiling is manually disabled.", - ) - - def test_enable_disable_slicing(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - - torch.manual_seed(0) - model = self.model_class(**init_dict).to(torch_device) - - inputs_dict.update({"return_dict": False}) - - torch.manual_seed(0) - output_without_slicing = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - torch.manual_seed(0) - model.enable_slicing() - output_with_slicing = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - self.assertLess( - (output_without_slicing.detach().cpu().numpy() - output_with_slicing.detach().cpu().numpy()).max(), - 0.5, - "VAE slicing should not affect the inference results", - ) - - torch.manual_seed(0) - model.disable_slicing() - output_without_slicing_2 = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - self.assertEqual( - output_without_slicing.detach().cpu().numpy().all(), - output_without_slicing_2.detach().cpu().numpy().all(), - "Without slicing outputs should match with the outputs when slicing is manually disabled.", - ) - def test_gradient_checkpointing_is_applied(self): expected_set = { "CogVideoXDownBlock3D", diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_temporal_decoder.py b/tests/models/autoencoders/test_models_autoencoder_kl_temporal_decoder.py index 6cb427bff8..93f40f44a9 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_temporal_decoder.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_temporal_decoder.py @@ -22,13 +22,14 @@ from ...testing_utils import ( floats_tensor, torch_device, ) -from ..test_modeling_common import ModelTesterMixin, UNetTesterMixin +from ..test_modeling_common import ModelTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() -class AutoencoderKLTemporalDecoderTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): +class AutoencoderKLTemporalDecoderTests(ModelTesterMixin, AutoencoderTesterMixin, unittest.TestCase): model_class = AutoencoderKLTemporalDecoder main_input_name = "sample" base_precision = 1e-2 @@ -67,7 +68,3 @@ class AutoencoderKLTemporalDecoderTests(ModelTesterMixin, UNetTesterMixin, unitt def test_gradient_checkpointing_is_applied(self): expected_set = {"Encoder", "TemporalDecoder", "UNetMidBlock2D"} super().test_gradient_checkpointing_is_applied(expected_set=expected_set) - - @unittest.skip("Test unsupported.") - def test_forward_with_norm_groups(self): - pass diff --git a/tests/models/autoencoders/test_models_autoencoder_ltx_video.py b/tests/models/autoencoders/test_models_autoencoder_ltx_video.py index 21ab3896c8..527be1b4ec 100644 --- a/tests/models/autoencoders/test_models_autoencoder_ltx_video.py +++ b/tests/models/autoencoders/test_models_autoencoder_ltx_video.py @@ -24,13 +24,14 @@ from ...testing_utils import ( floats_tensor, torch_device, ) -from ..test_modeling_common import ModelTesterMixin, UNetTesterMixin +from ..test_modeling_common import ModelTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() -class AutoencoderKLLTXVideo090Tests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): +class AutoencoderKLLTXVideo090Tests(ModelTesterMixin, AutoencoderTesterMixin, unittest.TestCase): model_class = AutoencoderKLLTXVideo main_input_name = "sample" base_precision = 1e-2 @@ -99,7 +100,7 @@ class AutoencoderKLLTXVideo090Tests(ModelTesterMixin, UNetTesterMixin, unittest. pass -class AutoencoderKLLTXVideo091Tests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): +class AutoencoderKLLTXVideo091Tests(ModelTesterMixin, unittest.TestCase): model_class = AutoencoderKLLTXVideo main_input_name = "sample" base_precision = 1e-2 @@ -167,34 +168,3 @@ class AutoencoderKLLTXVideo091Tests(ModelTesterMixin, UNetTesterMixin, unittest. @unittest.skip("AutoencoderKLLTXVideo does not support `norm_num_groups` because it does not use GroupNorm.") def test_forward_with_norm_groups(self): pass - - def test_enable_disable_tiling(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - - torch.manual_seed(0) - model = self.model_class(**init_dict).to(torch_device) - - inputs_dict.update({"return_dict": False}) - - torch.manual_seed(0) - output_without_tiling = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - torch.manual_seed(0) - model.enable_tiling() - output_with_tiling = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - self.assertLess( - (output_without_tiling.detach().cpu().numpy() - output_with_tiling.detach().cpu().numpy()).max(), - 0.5, - "VAE tiling should not affect the inference results", - ) - - torch.manual_seed(0) - model.disable_tiling() - output_without_tiling_2 = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - self.assertEqual( - output_without_tiling.detach().cpu().numpy().all(), - output_without_tiling_2.detach().cpu().numpy().all(), - "Without tiling outputs should match with the outputs when tiling is manually disabled.", - ) diff --git a/tests/models/autoencoders/test_models_autoencoder_magvit.py b/tests/models/autoencoders/test_models_autoencoder_magvit.py index 58cbfc05bd..f7304df140 100644 --- a/tests/models/autoencoders/test_models_autoencoder_magvit.py +++ b/tests/models/autoencoders/test_models_autoencoder_magvit.py @@ -18,13 +18,14 @@ import unittest from diffusers import AutoencoderKLMagvit from ...testing_utils import enable_full_determinism, floats_tensor, torch_device -from ..test_modeling_common import ModelTesterMixin, UNetTesterMixin +from ..test_modeling_common import ModelTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() -class AutoencoderKLMagvitTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): +class AutoencoderKLMagvitTests(ModelTesterMixin, AutoencoderTesterMixin, unittest.TestCase): model_class = AutoencoderKLMagvit main_input_name = "sample" base_precision = 1e-2 @@ -88,3 +89,9 @@ class AutoencoderKLMagvitTests(ModelTesterMixin, UNetTesterMixin, unittest.TestC @unittest.skip("Unsupported test.") def test_forward_with_norm_groups(self): pass + + @unittest.skip( + "Unsupported test. Error: RuntimeError: Sizes of tensors must match except in dimension 0. Expected size 9 but got size 12 for tensor number 1 in the list." + ) + def test_enable_disable_slicing(self): + pass diff --git a/tests/models/autoencoders/test_models_autoencoder_mochi.py b/tests/models/autoencoders/test_models_autoencoder_mochi.py index b8c5aaaa1e..ab8d429a67 100755 --- a/tests/models/autoencoders/test_models_autoencoder_mochi.py +++ b/tests/models/autoencoders/test_models_autoencoder_mochi.py @@ -17,18 +17,15 @@ import unittest from diffusers import AutoencoderKLMochi -from ...testing_utils import ( - enable_full_determinism, - floats_tensor, - torch_device, -) -from ..test_modeling_common import ModelTesterMixin, UNetTesterMixin +from ...testing_utils import enable_full_determinism, floats_tensor, torch_device +from ..test_modeling_common import ModelTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() -class AutoencoderKLMochiTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): +class AutoencoderKLMochiTests(ModelTesterMixin, AutoencoderTesterMixin, unittest.TestCase): model_class = AutoencoderKLMochi main_input_name = "sample" base_precision = 1e-2 @@ -79,14 +76,6 @@ class AutoencoderKLMochiTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCa } super().test_gradient_checkpointing_is_applied(expected_set=expected_set) - @unittest.skip("Unsupported test.") - def test_forward_with_norm_groups(self): - """ - tests/models/autoencoders/test_models_autoencoder_mochi.py::AutoencoderKLMochiTests::test_forward_with_norm_groups - - TypeError: AutoencoderKLMochi.__init__() got an unexpected keyword argument 'norm_num_groups' - """ - pass - @unittest.skip("Unsupported test.") def test_model_parallelism(self): """ diff --git a/tests/models/autoencoders/test_models_autoencoder_oobleck.py b/tests/models/autoencoders/test_models_autoencoder_oobleck.py index eb7bd50f4a..d10e8ba33a 100644 --- a/tests/models/autoencoders/test_models_autoencoder_oobleck.py +++ b/tests/models/autoencoders/test_models_autoencoder_oobleck.py @@ -30,13 +30,14 @@ from ...testing_utils import ( torch_all_close, torch_device, ) -from ..test_modeling_common import ModelTesterMixin, UNetTesterMixin +from ..test_modeling_common import ModelTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() -class AutoencoderOobleckTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): +class AutoencoderOobleckTests(ModelTesterMixin, AutoencoderTesterMixin, unittest.TestCase): model_class = AutoencoderOobleck main_input_name = "sample" base_precision = 1e-2 @@ -106,10 +107,6 @@ class AutoencoderOobleckTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCa "Without slicing outputs should match with the outputs when slicing is manually disabled.", ) - @unittest.skip("Test unsupported.") - def test_forward_with_norm_groups(self): - pass - @unittest.skip("No attention module used in this model") def test_set_attn_processor_for_determinism(self): return diff --git a/tests/models/autoencoders/test_models_autoencoder_tiny.py b/tests/models/autoencoders/test_models_autoencoder_tiny.py index 4d1dc69cfa..68232aa12f 100644 --- a/tests/models/autoencoders/test_models_autoencoder_tiny.py +++ b/tests/models/autoencoders/test_models_autoencoder_tiny.py @@ -31,13 +31,14 @@ from ...testing_utils import ( torch_all_close, torch_device, ) -from ..test_modeling_common import ModelTesterMixin, UNetTesterMixin +from ..test_modeling_common import ModelTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() -class AutoencoderTinyTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): +class AutoencoderTinyTests(ModelTesterMixin, AutoencoderTesterMixin, unittest.TestCase): model_class = AutoencoderTiny main_input_name = "sample" base_precision = 1e-2 @@ -81,37 +82,6 @@ class AutoencoderTinyTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase) def test_enable_disable_tiling(self): pass - def test_enable_disable_slicing(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - - torch.manual_seed(0) - model = self.model_class(**init_dict).to(torch_device) - - inputs_dict.update({"return_dict": False}) - - torch.manual_seed(0) - output_without_slicing = model(**inputs_dict)[0] - - torch.manual_seed(0) - model.enable_slicing() - output_with_slicing = model(**inputs_dict)[0] - - self.assertLess( - (output_without_slicing.detach().cpu().numpy() - output_with_slicing.detach().cpu().numpy()).max(), - 0.5, - "VAE slicing should not affect the inference results", - ) - - torch.manual_seed(0) - model.disable_slicing() - output_without_slicing_2 = model(**inputs_dict)[0] - - self.assertEqual( - output_without_slicing.detach().cpu().numpy().all(), - output_without_slicing_2.detach().cpu().numpy().all(), - "Without slicing outputs should match with the outputs when slicing is manually disabled.", - ) - @unittest.skip("Test not supported.") def test_outputs_equivalence(self): pass diff --git a/tests/models/autoencoders/test_models_autoencoder_wan.py b/tests/models/autoencoders/test_models_autoencoder_wan.py index cc9c888681..051098dc7a 100644 --- a/tests/models/autoencoders/test_models_autoencoder_wan.py +++ b/tests/models/autoencoders/test_models_autoencoder_wan.py @@ -15,18 +15,17 @@ import unittest -import torch - from diffusers import AutoencoderKLWan from ...testing_utils import enable_full_determinism, floats_tensor, torch_device -from ..test_modeling_common import ModelTesterMixin, UNetTesterMixin +from ..test_modeling_common import ModelTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() -class AutoencoderKLWanTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): +class AutoencoderKLWanTests(ModelTesterMixin, AutoencoderTesterMixin, unittest.TestCase): model_class = AutoencoderKLWan main_input_name = "sample" base_precision = 1e-2 @@ -76,68 +75,6 @@ class AutoencoderKLWanTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase inputs_dict = self.dummy_input_tiling return init_dict, inputs_dict - def test_enable_disable_tiling(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_tiling() - - torch.manual_seed(0) - model = self.model_class(**init_dict).to(torch_device) - - inputs_dict.update({"return_dict": False}) - - torch.manual_seed(0) - output_without_tiling = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - torch.manual_seed(0) - model.enable_tiling(96, 96, 64, 64) - output_with_tiling = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - self.assertLess( - (output_without_tiling.detach().cpu().numpy() - output_with_tiling.detach().cpu().numpy()).max(), - 0.5, - "VAE tiling should not affect the inference results", - ) - - torch.manual_seed(0) - model.disable_tiling() - output_without_tiling_2 = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - self.assertEqual( - output_without_tiling.detach().cpu().numpy().all(), - output_without_tiling_2.detach().cpu().numpy().all(), - "Without tiling outputs should match with the outputs when tiling is manually disabled.", - ) - - def test_enable_disable_slicing(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - - torch.manual_seed(0) - model = self.model_class(**init_dict).to(torch_device) - - inputs_dict.update({"return_dict": False}) - - torch.manual_seed(0) - output_without_slicing = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - torch.manual_seed(0) - model.enable_slicing() - output_with_slicing = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - self.assertLess( - (output_without_slicing.detach().cpu().numpy() - output_with_slicing.detach().cpu().numpy()).max(), - 0.05, - "VAE slicing should not affect the inference results", - ) - - torch.manual_seed(0) - model.disable_slicing() - output_without_slicing_2 = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - self.assertEqual( - output_without_slicing.detach().cpu().numpy().all(), - output_without_slicing_2.detach().cpu().numpy().all(), - "Without slicing outputs should match with the outputs when slicing is manually disabled.", - ) - @unittest.skip("Gradient checkpointing has not been implemented yet") def test_gradient_checkpointing_is_applied(self): pass diff --git a/tests/models/autoencoders/test_models_consistency_decoder_vae.py b/tests/models/autoencoders/test_models_consistency_decoder_vae.py index 7e44edba36..ef04d151ec 100644 --- a/tests/models/autoencoders/test_models_consistency_decoder_vae.py +++ b/tests/models/autoencoders/test_models_consistency_decoder_vae.py @@ -31,12 +31,13 @@ from ...testing_utils import ( torch_device, ) from ..test_modeling_common import ModelTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() -class ConsistencyDecoderVAETests(ModelTesterMixin, unittest.TestCase): +class ConsistencyDecoderVAETests(ModelTesterMixin, AutoencoderTesterMixin, unittest.TestCase): model_class = ConsistencyDecoderVAE main_input_name = "sample" base_precision = 1e-2 @@ -92,70 +93,6 @@ class ConsistencyDecoderVAETests(ModelTesterMixin, unittest.TestCase): def prepare_init_args_and_inputs_for_common(self): return self.init_dict, self.inputs_dict() - def test_enable_disable_tiling(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - - torch.manual_seed(0) - model = self.model_class(**init_dict).to(torch_device) - - inputs_dict.update({"return_dict": False}) - _ = inputs_dict.pop("generator") - - torch.manual_seed(0) - output_without_tiling = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - torch.manual_seed(0) - model.enable_tiling() - output_with_tiling = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - self.assertLess( - (output_without_tiling.detach().cpu().numpy() - output_with_tiling.detach().cpu().numpy()).max(), - 0.5, - "VAE tiling should not affect the inference results", - ) - - torch.manual_seed(0) - model.disable_tiling() - output_without_tiling_2 = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - self.assertEqual( - output_without_tiling.detach().cpu().numpy().all(), - output_without_tiling_2.detach().cpu().numpy().all(), - "Without tiling outputs should match with the outputs when tiling is manually disabled.", - ) - - def test_enable_disable_slicing(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - - torch.manual_seed(0) - model = self.model_class(**init_dict).to(torch_device) - - inputs_dict.update({"return_dict": False}) - _ = inputs_dict.pop("generator") - - torch.manual_seed(0) - output_without_slicing = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - torch.manual_seed(0) - model.enable_slicing() - output_with_slicing = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - self.assertLess( - (output_without_slicing.detach().cpu().numpy() - output_with_slicing.detach().cpu().numpy()).max(), - 0.5, - "VAE slicing should not affect the inference results", - ) - - torch.manual_seed(0) - model.disable_slicing() - output_without_slicing_2 = model(**inputs_dict, generator=torch.manual_seed(0))[0] - - self.assertEqual( - output_without_slicing.detach().cpu().numpy().all(), - output_without_slicing_2.detach().cpu().numpy().all(), - "Without slicing outputs should match with the outputs when slicing is manually disabled.", - ) - @slow class ConsistencyDecoderVAEIntegrationTests(unittest.TestCase): diff --git a/tests/models/autoencoders/test_models_vq.py b/tests/models/autoencoders/test_models_vq.py index 1c636b0817..b88d24d1f2 100644 --- a/tests/models/autoencoders/test_models_vq.py +++ b/tests/models/autoencoders/test_models_vq.py @@ -19,19 +19,15 @@ import torch from diffusers import VQModel -from ...testing_utils import ( - backend_manual_seed, - enable_full_determinism, - floats_tensor, - torch_device, -) -from ..test_modeling_common import ModelTesterMixin, UNetTesterMixin +from ...testing_utils import backend_manual_seed, enable_full_determinism, floats_tensor, torch_device +from ..test_modeling_common import ModelTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() -class VQModelTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): +class VQModelTests(ModelTesterMixin, AutoencoderTesterMixin, unittest.TestCase): model_class = VQModel main_input_name = "sample" diff --git a/tests/models/autoencoders/testing_utils.py b/tests/models/autoencoders/testing_utils.py new file mode 100644 index 0000000000..cf1f10a4a5 --- /dev/null +++ b/tests/models/autoencoders/testing_utils.py @@ -0,0 +1,142 @@ +import inspect + +import numpy as np +import pytest +import torch + +from diffusers.models.autoencoders.vae import DecoderOutput +from diffusers.utils.torch_utils import torch_device + + +class AutoencoderTesterMixin: + """ + Test mixin class specific to VAEs to test for slicing and tiling. Diffusion networks + usually don't do slicing and tiling. + """ + + @staticmethod + def _accepts_generator(model): + model_sig = inspect.signature(model.forward) + accepts_generator = "generator" in model_sig.parameters + return accepts_generator + + @staticmethod + def _accepts_norm_num_groups(model_class): + model_sig = inspect.signature(model_class.__init__) + accepts_norm_groups = "norm_num_groups" in model_sig.parameters + return accepts_norm_groups + + def test_forward_with_norm_groups(self): + if not self._accepts_norm_num_groups(self.model_class): + pytest.skip(f"Test not supported for {self.model_class.__name__}") + init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + + init_dict["norm_num_groups"] = 16 + init_dict["block_out_channels"] = (16, 32) + + model = self.model_class(**init_dict) + model.to(torch_device) + model.eval() + + with torch.no_grad(): + output = model(**inputs_dict) + + if isinstance(output, dict): + output = output.to_tuple()[0] + + self.assertIsNotNone(output) + expected_shape = inputs_dict["sample"].shape + self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match") + + def test_enable_disable_tiling(self): + if not hasattr(self.model_class, "enable_tiling"): + pytest.skip(f"Skipping test as {self.model_class.__name__} doesn't support tiling.") + + init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + + torch.manual_seed(0) + model = self.model_class(**init_dict).to(torch_device) + + inputs_dict.update({"return_dict": False}) + _ = inputs_dict.pop("generator", None) + accepts_generator = self._accepts_generator(model) + + torch.manual_seed(0) + if accepts_generator: + inputs_dict["generator"] = torch.manual_seed(0) + output_without_tiling = model(**inputs_dict)[0] + # Mochi-1 + if isinstance(output_without_tiling, DecoderOutput): + output_without_tiling = output_without_tiling.sample + + torch.manual_seed(0) + model.enable_tiling() + if accepts_generator: + inputs_dict["generator"] = torch.manual_seed(0) + output_with_tiling = model(**inputs_dict)[0] + if isinstance(output_with_tiling, DecoderOutput): + output_with_tiling = output_with_tiling.sample + + assert ( + output_without_tiling.detach().cpu().numpy() - output_with_tiling.detach().cpu().numpy() + ).max() < 0.5, "VAE tiling should not affect the inference results" + + torch.manual_seed(0) + model.disable_tiling() + if accepts_generator: + inputs_dict["generator"] = torch.manual_seed(0) + output_without_tiling_2 = model(**inputs_dict)[0] + if isinstance(output_without_tiling_2, DecoderOutput): + output_without_tiling_2 = output_without_tiling_2.sample + + assert np.allclose( + output_without_tiling.detach().cpu().numpy().all(), + output_without_tiling_2.detach().cpu().numpy().all(), + ), "Without tiling outputs should match with the outputs when tiling is manually disabled." + + def test_enable_disable_slicing(self): + if not hasattr(self.model_class, "enable_slicing"): + pytest.skip(f"Skipping test as {self.model_class.__name__} doesn't support slicing.") + + init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + + torch.manual_seed(0) + model = self.model_class(**init_dict).to(torch_device) + + inputs_dict.update({"return_dict": False}) + _ = inputs_dict.pop("generator", None) + accepts_generator = self._accepts_generator(model) + + if accepts_generator: + inputs_dict["generator"] = torch.manual_seed(0) + + torch.manual_seed(0) + output_without_slicing = model(**inputs_dict)[0] + # Mochi-1 + if isinstance(output_without_slicing, DecoderOutput): + output_without_slicing = output_without_slicing.sample + + torch.manual_seed(0) + model.enable_slicing() + if accepts_generator: + inputs_dict["generator"] = torch.manual_seed(0) + output_with_slicing = model(**inputs_dict)[0] + if isinstance(output_with_slicing, DecoderOutput): + output_with_slicing = output_with_slicing.sample + + assert ( + output_without_slicing.detach().cpu().numpy() - output_with_slicing.detach().cpu().numpy() + ).max() < 0.5, "VAE slicing should not affect the inference results" + + torch.manual_seed(0) + model.disable_slicing() + if accepts_generator: + inputs_dict["generator"] = torch.manual_seed(0) + output_without_slicing_2 = model(**inputs_dict)[0] + if isinstance(output_without_slicing_2, DecoderOutput): + output_without_slicing_2 = output_without_slicing_2.sample + + assert np.allclose( + output_without_slicing.detach().cpu().numpy().all(), + output_without_slicing_2.detach().cpu().numpy().all(), + ), "Without slicing outputs should match with the outputs when slicing is manually disabled." diff --git a/tests/models/test_modeling_common.py b/tests/models/test_modeling_common.py index a44ef571c5..6f4c3d544b 100644 --- a/tests/models/test_modeling_common.py +++ b/tests/models/test_modeling_common.py @@ -450,7 +450,15 @@ class ModelUtilsTest(unittest.TestCase): class UNetTesterMixin: + @staticmethod + def _accepts_norm_num_groups(model_class): + model_sig = inspect.signature(model_class.__init__) + accepts_norm_groups = "norm_num_groups" in model_sig.parameters + return accepts_norm_groups + def test_forward_with_norm_groups(self): + if not self._accepts_norm_num_groups(self.model_class): + pytest.skip(f"Test not supported for {self.model_class.__name__}") init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() init_dict["norm_num_groups"] = 16 From 1b456bd5d52faa0896d9ebceed38e1d4f298a56f Mon Sep 17 00:00:00 2001 From: Ali Imran <80049236+EazyAl@users.noreply.github.com> Date: Fri, 17 Oct 2025 23:10:50 +0200 Subject: [PATCH 073/123] docs: cleanup of runway model (#12503) * cleanup of runway model * quality fixes --- README.md | 2 +- .../en/api/models/asymmetricautoencoderkl.md | 2 +- .../pipelines/stable_diffusion/depth2img.md | 2 +- .../api/pipelines/stable_diffusion/inpaint.md | 4 +-- .../stable_diffusion/latent_upscale.md | 2 +- .../pipelines/stable_diffusion/overview.md | 4 +-- .../stable_diffusion/stable_diffusion_2.md | 2 +- .../pipelines/stable_diffusion/text2img.md | 2 +- .../api/pipelines/stable_diffusion/upscale.md | 2 +- docs/source/en/training/adapt_a_model.md | 4 +-- .../en/using-diffusers/inference_with_lcm.md | 2 +- docs/source/en/using-diffusers/inpaint.md | 30 +++++++++---------- docs/source/ko/optimization/torch2.0.md | 2 +- docs/source/ko/training/adapt_a_model.md | 4 +-- docs/source/ko/using-diffusers/inpaint.md | 13 ++------ docs/source/zh/training/adapt_a_model.md | 4 +-- examples/community/README.md | 2 +- .../community/adaptive_mask_inpainting.py | 12 ++++---- .../community/composable_stable_diffusion.py | 6 ++-- examples/community/fresco_v2v.py | 2 +- examples/community/hd_painter.py | 2 +- examples/community/img2img_inpainting.py | 2 +- examples/community/instaflow_one_step.py | 6 ++-- examples/community/ip_adapter_face_id.py | 6 ++-- examples/community/kohya_hires_fix.py | 2 +- .../latent_consistency_interpolate.py | 2 +- examples/community/llm_grounded_diffusion.py | 6 ++-- examples/community/lpw_stable_diffusion.py | 4 +-- examples/community/lpw_stable_diffusion_xl.py | 2 +- examples/community/matryoshka.py | 4 +-- .../multilingual_stable_diffusion.py | 2 +- .../pipeline_controlnet_xl_kolors_inpaint.py | 2 +- examples/community/pipeline_fabric.py | 6 ++-- .../community/pipeline_kolors_inpainting.py | 2 +- examples/community/pipeline_prompt2prompt.py | 6 ++-- .../community/pipeline_sdxl_style_aligned.py | 2 +- .../pipeline_stable_diffusion_boxdiff.py | 8 ++--- .../pipeline_stable_diffusion_pag.py | 8 ++--- ...pipeline_stable_diffusion_upscale_ldm3d.py | 2 +- ...ne_stable_diffusion_xl_attentive_eraser.py | 2 +- ..._stable_diffusion_xl_controlnet_adapter.py | 2 +- ...diffusion_xl_controlnet_adapter_inpaint.py | 4 +-- examples/community/pipeline_zero1to3.py | 8 ++--- examples/community/rerender_a_video.py | 2 +- examples/community/run_onnx_controlnet.py | 2 +- examples/community/run_tensorrt_controlnet.py | 2 +- examples/community/sd_text2img_k_diffusion.py | 2 +- examples/community/sde_drag.py | 2 +- .../community/stable_diffusion_comparison.py | 2 +- .../stable_diffusion_controlnet_img2img.py | 2 +- .../stable_diffusion_controlnet_inpaint.py | 2 +- ...le_diffusion_controlnet_inpaint_img2img.py | 2 +- .../stable_diffusion_controlnet_reference.py | 2 +- examples/community/stable_diffusion_ipex.py | 8 ++--- examples/community/stable_diffusion_mega.py | 2 +- .../community/stable_diffusion_reference.py | 8 ++--- .../community/stable_diffusion_repaint.py | 6 ++-- .../stable_diffusion_tensorrt_img2img.py | 6 ++-- .../stable_diffusion_tensorrt_inpaint.py | 6 ++-- .../stable_diffusion_tensorrt_txt2img.py | 6 ++-- examples/community/text_inpainting.py | 2 +- examples/research_projects/anytext/anytext.py | 2 +- .../dreambooth_inpaint/README.md | 8 ++--- .../research_projects/ip_adapter/README.md | 4 +-- .../README.md | 2 +- .../pipeline_prompt_diffusion.py | 2 +- .../research_projects/vae/vae_roundtrip.py | 2 +- scripts/convert_consistency_decoder.py | 11 ++++--- src/diffusers/models/auto_model.py | 4 +-- src/diffusers/models/modeling_flax_utils.py | 18 +++++------ src/diffusers/models/modeling_utils.py | 6 ++-- .../stable_diffusion_xl/denoise.py | 2 +- src/diffusers/pipelines/README.md | 2 +- .../controlnet_xs/pipeline_controlnet_xs.py | 4 +-- .../alt_diffusion/pipeline_alt_diffusion.py | 8 ++--- .../pipeline_alt_diffusion_img2img.py | 8 ++--- .../pipeline_cycle_diffusion.py | 8 ++--- ...ne_onnx_stable_diffusion_inpaint_legacy.py | 3 +- ...ipeline_stable_diffusion_inpaint_legacy.py | 7 +++-- ...pipeline_stable_diffusion_model_editing.py | 4 +-- .../pipeline_stable_diffusion_paradigms.py | 10 ++++--- .../pipeline_stable_diffusion_pix2pix_zero.py | 3 +- .../pipeline_versatile_diffusion.py | 4 +-- .../pipeline_latent_consistency_img2img.py | 4 +-- .../pipeline_latent_consistency_text2img.py | 4 +-- .../pipeline_leditspp_stable_diffusion.py | 6 ++-- .../pag/pipeline_pag_controlnet_sd.py | 9 ++++-- .../pag/pipeline_pag_controlnet_sd_inpaint.py | 20 ++++++++----- .../pipelines/pag/pipeline_pag_sd.py | 10 +++---- .../pipelines/pag/pipeline_pag_sd_img2img.py | 10 +++---- .../pipelines/pag/pipeline_pag_sd_inpaint.py | 12 ++++---- .../pag/pipeline_pag_sd_xl_inpaint.py | 2 +- .../pipeline_paint_by_example.py | 4 +-- .../pipeline_semantic_stable_diffusion.py | 6 ++-- .../pipeline_stable_diffusion_3_inpaint.py | 2 +- .../pipeline_stable_diffusion_xl_inpaint.py | 2 +- src/diffusers/utils/hub_utils.py | 3 +- tests/pipelines/pag/test_pag_sd_inpaint.py | 2 +- 98 files changed, 241 insertions(+), 234 deletions(-) diff --git a/README.md b/README.md index 68202ba095..07334416ed 100644 --- a/README.md +++ b/README.md @@ -171,7 +171,7 @@ Also, say 👋 in our public Discord channel Text-guided Image Inpainting Stable Diffusion Inpainting - runwayml/stable-diffusion-inpainting + stable-diffusion-v1-5/stable-diffusion-inpainting Image Variation diff --git a/docs/source/en/api/models/asymmetricautoencoderkl.md b/docs/source/en/api/models/asymmetricautoencoderkl.md index fbadf9bd40..0c7fbd38fb 100644 --- a/docs/source/en/api/models/asymmetricautoencoderkl.md +++ b/docs/source/en/api/models/asymmetricautoencoderkl.md @@ -39,7 +39,7 @@ mask_url = "https://huggingface.co/datasets/hf-internal-testing/diffusers-images original_image = load_image(img_url).resize((512, 512)) mask_image = load_image(mask_url).resize((512, 512)) -pipe = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-inpainting") +pipe = StableDiffusionInpaintPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-inpainting") pipe.vae = AsymmetricAutoencoderKL.from_pretrained("cross-attention/asymmetric-autoencoder-kl-x-1-5") pipe.to("cuda") diff --git a/docs/source/en/api/pipelines/stable_diffusion/depth2img.md b/docs/source/en/api/pipelines/stable_diffusion/depth2img.md index aa43cf7db9..9ac9b7f906 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/depth2img.md +++ b/docs/source/en/api/pipelines/stable_diffusion/depth2img.md @@ -21,7 +21,7 @@ The Stable Diffusion model can also infer depth based on an image using [MiDaS]( > [!TIP] > Make sure to check out the Stable Diffusion [Tips](overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality, and how to reuse pipeline components efficiently! > -> If you're interested in using one of the official checkpoints for a task, explore the [CompVis](https://huggingface.co/CompVis), [Runway](https://huggingface.co/runwayml), and [Stability AI](https://huggingface.co/stabilityai) Hub organizations! +> If you're interested in using one of the official checkpoints for a task, explore the [CompVis](https://huggingface.co/CompVis) and [Stability AI](https://huggingface.co/stabilityai) Hub organizations! ## StableDiffusionDepth2ImgPipeline diff --git a/docs/source/en/api/pipelines/stable_diffusion/inpaint.md b/docs/source/en/api/pipelines/stable_diffusion/inpaint.md index 84cc31e158..d60f472d0a 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/inpaint.md +++ b/docs/source/en/api/pipelines/stable_diffusion/inpaint.md @@ -21,14 +21,14 @@ The Stable Diffusion model can also be applied to inpainting which lets you edit ## Tips It is recommended to use this pipeline with checkpoints that have been specifically fine-tuned for inpainting, such -as [runwayml/stable-diffusion-inpainting](https://huggingface.co/runwayml/stable-diffusion-inpainting). Default +as [stable-diffusion-v1-5/stable-diffusion-inpainting](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-inpainting). Default text-to-image Stable Diffusion checkpoints, such as [stable-diffusion-v1-5/stable-diffusion-v1-5](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) are also compatible but they might be less performant. > [!TIP] > Make sure to check out the Stable Diffusion [Tips](overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality, and how to reuse pipeline components efficiently! > -> If you're interested in using one of the official checkpoints for a task, explore the [CompVis](https://huggingface.co/CompVis), [Runway](https://huggingface.co/runwayml), and [Stability AI](https://huggingface.co/stabilityai) Hub organizations! +> If you're interested in using one of the official checkpoints for a task, explore the [CompVis](https://huggingface.co/CompVis) and [Stability AI](https://huggingface.co/stabilityai) Hub organizations! ## StableDiffusionInpaintPipeline diff --git a/docs/source/en/api/pipelines/stable_diffusion/latent_upscale.md b/docs/source/en/api/pipelines/stable_diffusion/latent_upscale.md index 4f0521740c..19eae9a9ce 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/latent_upscale.md +++ b/docs/source/en/api/pipelines/stable_diffusion/latent_upscale.md @@ -17,7 +17,7 @@ The Stable Diffusion latent upscaler model was created by [Katherine Crowson](ht > [!TIP] > Make sure to check out the Stable Diffusion [Tips](overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality, and how to reuse pipeline components efficiently! > -> If you're interested in using one of the official checkpoints for a task, explore the [CompVis](https://huggingface.co/CompVis), [Runway](https://huggingface.co/runwayml), and [Stability AI](https://huggingface.co/stabilityai) Hub organizations! +> If you're interested in using one of the official checkpoints for a task, explore the [CompVis](https://huggingface.co/CompVis) and [Stability AI](https://huggingface.co/stabilityai) Hub organizations! ## StableDiffusionLatentUpscalePipeline diff --git a/docs/source/en/api/pipelines/stable_diffusion/overview.md b/docs/source/en/api/pipelines/stable_diffusion/overview.md index 7e6e16c347..2d2de39c91 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/overview.md +++ b/docs/source/en/api/pipelines/stable_diffusion/overview.md @@ -22,7 +22,7 @@ Stable Diffusion is trained on 512x512 images from a subset of the LAION-5B data For more details about how Stable Diffusion works and how it differs from the base latent diffusion model, take a look at the Stability AI [announcement](https://stability.ai/blog/stable-diffusion-announcement) and our own [blog post](https://huggingface.co/blog/stable_diffusion#how-does-stable-diffusion-work) for more technical details. -You can find the original codebase for Stable Diffusion v1.0 at [CompVis/stable-diffusion](https://github.com/CompVis/stable-diffusion) and Stable Diffusion v2.0 at [Stability-AI/stablediffusion](https://github.com/Stability-AI/stablediffusion) as well as their original scripts for various tasks. Additional official checkpoints for the different Stable Diffusion versions and tasks can be found on the [CompVis](https://huggingface.co/CompVis), [Runway](https://huggingface.co/runwayml), and [Stability AI](https://huggingface.co/stabilityai) Hub organizations. Explore these organizations to find the best checkpoint for your use-case! +You can find the original codebase for Stable Diffusion v1.0 at [CompVis/stable-diffusion](https://github.com/CompVis/stable-diffusion) and Stable Diffusion v2.0 at [Stability-AI/stablediffusion](https://github.com/Stability-AI/stablediffusion) as well as their original scripts for various tasks. Additional official checkpoints for the different Stable Diffusion versions and tasks can be found on the [CompVis](https://huggingface.co/CompVis) and [Stability AI](https://huggingface.co/stabilityai) Hub organizations. Explore these organizations to find the best checkpoint for your use-case! The table below summarizes the available Stable Diffusion pipelines, their supported tasks, and an interactive demo: @@ -64,7 +64,7 @@ The table below summarizes the available Stable Diffusion pipelines, their suppo StableDiffusionInpaint inpainting - + diff --git a/docs/source/en/api/pipelines/stable_diffusion/stable_diffusion_2.md b/docs/source/en/api/pipelines/stable_diffusion/stable_diffusion_2.md index 67729cd195..31164fb9ae 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/stable_diffusion_2.md +++ b/docs/source/en/api/pipelines/stable_diffusion/stable_diffusion_2.md @@ -36,7 +36,7 @@ Here are some examples for how to use Stable Diffusion 2 for each task: > [!TIP] > Make sure to check out the Stable Diffusion [Tips](overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality, and how to reuse pipeline components efficiently! > -> If you're interested in using one of the official checkpoints for a task, explore the [CompVis](https://huggingface.co/CompVis), [Runway](https://huggingface.co/runwayml), and [Stability AI](https://huggingface.co/stabilityai) Hub organizations! +> If you're interested in using one of the official checkpoints for a task, explore the [CompVis](https://huggingface.co/CompVis) and [Stability AI](https://huggingface.co/stabilityai) Hub organizations! ## Text-to-image diff --git a/docs/source/en/api/pipelines/stable_diffusion/text2img.md b/docs/source/en/api/pipelines/stable_diffusion/text2img.md index 59a0f00d22..2033638a1a 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/text2img.md +++ b/docs/source/en/api/pipelines/stable_diffusion/text2img.md @@ -25,7 +25,7 @@ The abstract from the paper is: > [!TIP] > Make sure to check out the Stable Diffusion [Tips](overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality, and how to reuse pipeline components efficiently! > -> If you're interested in using one of the official checkpoints for a task, explore the [CompVis](https://huggingface.co/CompVis), [Runway](https://huggingface.co/runwayml), and [Stability AI](https://huggingface.co/stabilityai) Hub organizations! +> If you're interested in using one of the official checkpoints for a task, explore the [CompVis](https://huggingface.co/CompVis) and [Stability AI](https://huggingface.co/stabilityai) Hub organizations! ## StableDiffusionPipeline diff --git a/docs/source/en/api/pipelines/stable_diffusion/upscale.md b/docs/source/en/api/pipelines/stable_diffusion/upscale.md index 14393370be..e7d571e73c 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/upscale.md +++ b/docs/source/en/api/pipelines/stable_diffusion/upscale.md @@ -21,7 +21,7 @@ The Stable Diffusion upscaler diffusion model was created by the researchers and > [!TIP] > Make sure to check out the Stable Diffusion [Tips](overview#tips) section to learn how to explore the tradeoff between scheduler speed and quality, and how to reuse pipeline components efficiently! > -> If you're interested in using one of the official checkpoints for a task, explore the [CompVis](https://huggingface.co/CompVis), [Runway](https://huggingface.co/runwayml), and [Stability AI](https://huggingface.co/stabilityai) Hub organizations! +> If you're interested in using one of the official checkpoints for a task, explore the [CompVis](https://huggingface.co/CompVis) and [Stability AI](https://huggingface.co/stabilityai) Hub organizations! ## StableDiffusionUpscalePipeline diff --git a/docs/source/en/training/adapt_a_model.md b/docs/source/en/training/adapt_a_model.md index f528c8bfb6..9b7efd2abf 100644 --- a/docs/source/en/training/adapt_a_model.md +++ b/docs/source/en/training/adapt_a_model.md @@ -16,12 +16,12 @@ pipeline.unet.config["in_channels"] 4 ``` -Inpainting requires 9 channels in the input sample. You can check this value in a pretrained inpainting model like [`runwayml/stable-diffusion-inpainting`](https://huggingface.co/runwayml/stable-diffusion-inpainting): +Inpainting requires 9 channels in the input sample. You can check this value in a pretrained inpainting model like [`stable-diffusion-v1-5/stable-diffusion-inpainting`](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-inpainting): ```py from diffusers import StableDiffusionPipeline -pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-inpainting", use_safetensors=True) +pipeline = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-inpainting", use_safetensors=True) pipeline.unet.config["in_channels"] 9 ``` diff --git a/docs/source/en/using-diffusers/inference_with_lcm.md b/docs/source/en/using-diffusers/inference_with_lcm.md index cde4168d38..258ed2979a 100644 --- a/docs/source/en/using-diffusers/inference_with_lcm.md +++ b/docs/source/en/using-diffusers/inference_with_lcm.md @@ -215,7 +215,7 @@ from diffusers import AutoPipelineForInpainting, LCMScheduler from diffusers.utils import load_image, make_image_grid pipe = AutoPipelineForInpainting.from_pretrained( - "runwayml/stable-diffusion-inpainting", + "stable-diffusion-v1-5/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16", ).to("cuda") diff --git a/docs/source/en/using-diffusers/inpaint.md b/docs/source/en/using-diffusers/inpaint.md index 28da3a68a5..232dbf2c6b 100644 --- a/docs/source/en/using-diffusers/inpaint.md +++ b/docs/source/en/using-diffusers/inpaint.md @@ -112,7 +112,7 @@ blurred_mask ## Popular models -[Stable Diffusion Inpainting](https://huggingface.co/runwayml/stable-diffusion-inpainting), [Stable Diffusion XL (SDXL) Inpainting](https://huggingface.co/diffusers/stable-diffusion-xl-1.0-inpainting-0.1), and [Kandinsky 2.2 Inpainting](https://huggingface.co/kandinsky-community/kandinsky-2-2-decoder-inpaint) are among the most popular models for inpainting. SDXL typically produces higher resolution images than Stable Diffusion v1.5, and Kandinsky 2.2 is also capable of generating high-quality images. +[Stable Diffusion Inpainting](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-inpainting), [Stable Diffusion XL (SDXL) Inpainting](https://huggingface.co/diffusers/stable-diffusion-xl-1.0-inpainting-0.1), and [Kandinsky 2.2 Inpainting](https://huggingface.co/kandinsky-community/kandinsky-2-2-decoder-inpaint) are among the most popular models for inpainting. SDXL typically produces higher resolution images than Stable Diffusion v1.5, and Kandinsky 2.2 is also capable of generating high-quality images. ### Stable Diffusion Inpainting @@ -124,7 +124,7 @@ from diffusers import AutoPipelineForInpainting from diffusers.utils import load_image, make_image_grid pipeline = AutoPipelineForInpainting.from_pretrained( - "runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16" + "stable-diffusion-v1-5/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16" ) pipeline.enable_model_cpu_offload() # remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed @@ -244,7 +244,7 @@ make_image_grid([init_image, image], rows=1, cols=2) ``` - + ```py import torch @@ -252,7 +252,7 @@ from diffusers import AutoPipelineForInpainting from diffusers.utils import load_image, make_image_grid pipeline = AutoPipelineForInpainting.from_pretrained( - "runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16" + "stable-diffusion-v1-5/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16" ) pipeline.enable_model_cpu_offload() # remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed @@ -278,7 +278,7 @@ make_image_grid([init_image, image], rows=1, cols=2)
-
runwayml/stable-diffusion-inpainting
+
stable-diffusion-v1-5/stable-diffusion-inpainting
@@ -308,7 +308,7 @@ make_image_grid([init_image, image], rows=1, cols=2) ``` - + ```py import torch @@ -316,7 +316,7 @@ from diffusers import AutoPipelineForInpainting from diffusers.utils import load_image, make_image_grid pipeline = AutoPipelineForInpainting.from_pretrained( - "runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16" + "stable-diffusion-v1-5/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16" ) pipeline.enable_model_cpu_offload() # remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed @@ -340,7 +340,7 @@ make_image_grid([init_image, image], rows=1, cols=2)
-
runwayml/stable-diffusion-inpainting
+
stable-diffusion-v1-5/stable-diffusion-inpainting
@@ -358,7 +358,7 @@ from diffusers.utils import load_image, make_image_grid device = "cuda" pipeline = AutoPipelineForInpainting.from_pretrained( - "runwayml/stable-diffusion-inpainting", + "stable-diffusion-v1-5/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16" ) @@ -396,7 +396,7 @@ from diffusers import AutoPipelineForInpainting from diffusers.utils import load_image, make_image_grid pipeline = AutoPipelineForInpainting.from_pretrained( - "runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16" + "stable-diffusion-v1-5/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16" ) pipeline.enable_model_cpu_offload() # remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed @@ -441,7 +441,7 @@ from diffusers import AutoPipelineForInpainting from diffusers.utils import load_image, make_image_grid pipeline = AutoPipelineForInpainting.from_pretrained( - "runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16" + "stable-diffusion-v1-5/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16" ) pipeline.enable_model_cpu_offload() # remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed @@ -481,7 +481,7 @@ from diffusers import AutoPipelineForInpainting from diffusers.utils import load_image, make_image_grid pipeline = AutoPipelineForInpainting.from_pretrained( - "runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16" + "stable-diffusion-v1-5/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16" ) pipeline.enable_model_cpu_offload() # remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed @@ -606,7 +606,7 @@ from diffusers import AutoPipelineForInpainting, AutoPipelineForImage2Image from diffusers.utils import load_image, make_image_grid pipeline = AutoPipelineForInpainting.from_pretrained( - "runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16" + "stable-diffusion-v1-5/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16" ) pipeline.enable_model_cpu_offload() # remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed @@ -683,7 +683,7 @@ from diffusers import AutoPipelineForInpainting from diffusers.utils import make_image_grid pipeline = AutoPipelineForInpainting.from_pretrained( - "runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, + "stable-diffusion-v1-5/stable-diffusion-inpainting", torch_dtype=torch.float16, ) pipeline.enable_model_cpu_offload() # remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed @@ -714,7 +714,7 @@ controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_inpai # pass ControlNet to the pipeline pipeline = StableDiffusionControlNetInpaintPipeline.from_pretrained( - "runwayml/stable-diffusion-inpainting", controlnet=controlnet, torch_dtype=torch.float16, variant="fp16" + "stable-diffusion-v1-5/stable-diffusion-inpainting", controlnet=controlnet, torch_dtype=torch.float16, variant="fp16" ) pipeline.enable_model_cpu_offload() # remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed diff --git a/docs/source/ko/optimization/torch2.0.md b/docs/source/ko/optimization/torch2.0.md index c78c4a87b6..354f7243cf 100644 --- a/docs/source/ko/optimization/torch2.0.md +++ b/docs/source/ko/optimization/torch2.0.md @@ -173,7 +173,7 @@ mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data init_image = download_image(img_url).resize((512, 512)) mask_image = download_image(mask_url).resize((512, 512)) -path = "runwayml/stable-diffusion-inpainting" +path = "stable-diffusion-v1-5/stable-diffusion-inpainting" run_compile = True # Set True / False diff --git a/docs/source/ko/training/adapt_a_model.md b/docs/source/ko/training/adapt_a_model.md index 3795558f5f..fe6fde05b7 100644 --- a/docs/source/ko/training/adapt_a_model.md +++ b/docs/source/ko/training/adapt_a_model.md @@ -28,12 +28,12 @@ pipeline.unet.config["in_channels"] 4 ``` -인페인팅은 입력 샘플에 9개의 채널이 필요합니다. [`runwayml/stable-diffusion-inpainting`](https://huggingface.co/runwayml/stable-diffusion-inpainting)와 같은 사전학습된 인페인팅 모델에서 이 값을 확인할 수 있습니다: +인페인팅은 입력 샘플에 9개의 채널이 필요합니다. [`stable-diffusion-v1-5/stable-diffusion-inpainting`](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-inpainting)와 같은 사전학습된 인페인팅 모델에서 이 값을 확인할 수 있습니다: ```py from diffusers import StableDiffusionPipeline -pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-inpainting") +pipeline = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-inpainting") pipeline.unet.config["in_channels"] 9 ``` diff --git a/docs/source/ko/using-diffusers/inpaint.md b/docs/source/ko/using-diffusers/inpaint.md index cefb892186..6c0c08bf73 100644 --- a/docs/source/ko/using-diffusers/inpaint.md +++ b/docs/source/ko/using-diffusers/inpaint.md @@ -14,7 +14,7 @@ specific language governing permissions and limitations under the License. [[open-in-colab]] -[`StableDiffusionInpaintPipeline`]은 마스크와 텍스트 프롬프트를 제공하여 이미지의 특정 부분을 편집할 수 있도록 합니다. 이 기능은 인페인팅 작업을 위해 특별히 훈련된 [`runwayml/stable-diffusion-inpainting`](https://huggingface.co/runwayml/stable-diffusion-inpainting)과 같은 Stable Diffusion 버전을 사용합니다. +[`StableDiffusionInpaintPipeline`]은 마스크와 텍스트 프롬프트를 제공하여 이미지의 특정 부분을 편집할 수 있도록 합니다. 이 기능은 인페인팅 작업을 위해 특별히 훈련된 [`stable-diffusion-v1-5/stable-diffusion-inpainting`](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-inpainting)과 같은 Stable Diffusion 버전을 사용합니다. 먼저 [`StableDiffusionInpaintPipeline`] 인스턴스를 불러옵니다: @@ -27,7 +27,7 @@ from io import BytesIO from diffusers import StableDiffusionInpaintPipeline pipeline = StableDiffusionInpaintPipeline.from_pretrained( - "runwayml/stable-diffusion-inpainting", + "stable-diffusion-v1-5/stable-diffusion-inpainting", torch_dtype=torch.float16, ) pipeline = pipeline.to("cuda") @@ -61,12 +61,3 @@ image = pipe(prompt=prompt, image=init_image, mask_image=mask_image).images[0] > [!WARNING] > 이전의 실험적인 인페인팅 구현에서는 품질이 낮은 다른 프로세스를 사용했습니다. 이전 버전과의 호환성을 보장하기 위해 새 모델이 포함되지 않은 사전학습된 파이프라인을 불러오면 이전 인페인팅 방법이 계속 적용됩니다. - -아래 Space에서 이미지 인페인팅을 직접 해보세요! - - diff --git a/docs/source/zh/training/adapt_a_model.md b/docs/source/zh/training/adapt_a_model.md index b5f9155697..7dbf46ec12 100644 --- a/docs/source/zh/training/adapt_a_model.md +++ b/docs/source/zh/training/adapt_a_model.md @@ -16,12 +16,12 @@ pipeline.unet.config["in_channels"] 4 ``` -而图像修复任务需要输入样本具有9个通道。您可以在 [`runwayml/stable-diffusion-inpainting`](https://huggingface.co/runwayml/stable-diffusion-inpainting) 这样的预训练修复模型中验证此参数： +而图像修复任务需要输入样本具有9个通道。您可以在 [`stable-diffusion-v1-5/stable-diffusion-inpainting`](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-inpainting) 这样的预训练修复模型中验证此参数： ```python from diffusers import StableDiffusionPipeline -pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-inpainting", use_safetensors=True) +pipeline = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-inpainting", use_safetensors=True) pipeline.unet.config["in_channels"] 9 ``` diff --git a/examples/community/README.md b/examples/community/README.md index e314463077..4a4b0f5fd9 100644 --- a/examples/community/README.md +++ b/examples/community/README.md @@ -1328,7 +1328,7 @@ model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined" # Load Stable Diffusion Inpainting Pipeline with custom pipeline pipe = DiffusionPipeline.from_pretrained( - "runwayml/stable-diffusion-inpainting", + "stable-diffusion-v1-5/stable-diffusion-inpainting", custom_pipeline="text_inpainting", segmentation_model=model, segmentation_processor=processor diff --git a/examples/community/adaptive_mask_inpainting.py b/examples/community/adaptive_mask_inpainting.py index aac460cb46..da67debe72 100644 --- a/examples/community/adaptive_mask_inpainting.py +++ b/examples/community/adaptive_mask_inpainting.py @@ -126,7 +126,7 @@ EXAMPLE_DOC_STRING = """ ... "lllyasviel/control_v11p_sd15_inpaint", torch_dtype=torch.float16 ... ) >>> pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained( - ... "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16 + ... "stable-diffusion-v1-5/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16 ... ) >>> pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config) @@ -347,7 +347,7 @@ class AdaptiveMaskInpaintPipeline( [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. @@ -429,8 +429,8 @@ class AdaptiveMaskInpaintPipeline( "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely .If you're checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" + " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" @@ -970,7 +970,7 @@ class AdaptiveMaskInpaintPipeline( >>> default_mask_image = download_image(mask_url).resize((512, 512)) >>> pipe = AdaptiveMaskInpaintPipeline.from_pretrained( - ... "runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16 + ... "stable-diffusion-v1-5/stable-diffusion-inpainting", torch_dtype=torch.float16 ... ) >>> pipe = pipe.to("cuda") @@ -1095,7 +1095,7 @@ class AdaptiveMaskInpaintPipeline( # 8. Check that sizes of mask, masked image and latents match if num_channels_unet == 9: - # default case for runwayml/stable-diffusion-inpainting + # default case for stable-diffusion-v1-5/stable-diffusion-inpainting num_channels_mask = mask.shape[1] num_channels_masked_image = masked_image_latents.shape[1] if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels: diff --git a/examples/community/composable_stable_diffusion.py b/examples/community/composable_stable_diffusion.py index a7c540ceb9..a7c610ad43 100644 --- a/examples/community/composable_stable_diffusion.py +++ b/examples/community/composable_stable_diffusion.py @@ -62,7 +62,7 @@ class ComposableStableDiffusionPipeline(DiffusionPipeline, StableDiffusionMixin) [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details. + Please, refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for details. feature_extractor ([`CLIPImageProcessor`]): Model that extracts features from generated images to be used as inputs for the `safety_checker`. """ @@ -145,8 +145,8 @@ class ComposableStableDiffusionPipeline(DiffusionPipeline, StableDiffusionMixin) "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" + " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" diff --git a/examples/community/fresco_v2v.py b/examples/community/fresco_v2v.py index 47ba71299d..b79834db5e 100644 --- a/examples/community/fresco_v2v.py +++ b/examples/community/fresco_v2v.py @@ -1276,7 +1276,7 @@ class FrescoV2VPipeline(StableDiffusionControlNetImg2ImgPipeline): [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. diff --git a/examples/community/hd_painter.py b/examples/community/hd_painter.py index 20bb43a76f..70e5656855 100644 --- a/examples/community/hd_painter.py +++ b/examples/community/hd_painter.py @@ -678,7 +678,7 @@ class StableDiffusionHDPainterPipeline(StableDiffusionInpaintPipeline): # 8. Check that sizes of mask, masked image and latents match if num_channels_unet == 9: - # default case for runwayml/stable-diffusion-inpainting + # default case for stable-diffusion-v1-5/stable-diffusion-inpainting num_channels_mask = mask.shape[1] num_channels_masked_image = masked_image_latents.shape[1] if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels: diff --git a/examples/community/img2img_inpainting.py b/examples/community/img2img_inpainting.py index 499230b1e2..595df107ca 100644 --- a/examples/community/img2img_inpainting.py +++ b/examples/community/img2img_inpainting.py @@ -78,7 +78,7 @@ class ImageToImageInpaintingPipeline(DiffusionPipeline): [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details. + Please, refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for details. feature_extractor ([`CLIPImageProcessor`]): Model that extracts features from generated images to be used as inputs for the `safety_checker`. """ diff --git a/examples/community/instaflow_one_step.py b/examples/community/instaflow_one_step.py index 06be1d10b6..0f16707ead 100644 --- a/examples/community/instaflow_one_step.py +++ b/examples/community/instaflow_one_step.py @@ -86,7 +86,7 @@ class InstaFlowPipeline( [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. @@ -165,8 +165,8 @@ class InstaFlowPipeline( "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" + " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" diff --git a/examples/community/ip_adapter_face_id.py b/examples/community/ip_adapter_face_id.py index 5b420882e9..d16aaf5a54 100644 --- a/examples/community/ip_adapter_face_id.py +++ b/examples/community/ip_adapter_face_id.py @@ -166,7 +166,7 @@ class IPAdapterFaceIDStableDiffusionPipeline( [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. @@ -247,8 +247,8 @@ class IPAdapterFaceIDStableDiffusionPipeline( "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" + " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" diff --git a/examples/community/kohya_hires_fix.py b/examples/community/kohya_hires_fix.py index 63f6b8973c..c968ecf2af 100644 --- a/examples/community/kohya_hires_fix.py +++ b/examples/community/kohya_hires_fix.py @@ -414,7 +414,7 @@ class StableDiffusionHighResFixPipeline(StableDiffusionPipeline): [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. diff --git a/examples/community/latent_consistency_interpolate.py b/examples/community/latent_consistency_interpolate.py index 9fc4233682..e8349ba317 100644 --- a/examples/community/latent_consistency_interpolate.py +++ b/examples/community/latent_consistency_interpolate.py @@ -222,7 +222,7 @@ class LatentConsistencyModelWalkPipeline( supports [`LCMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. diff --git a/examples/community/llm_grounded_diffusion.py b/examples/community/llm_grounded_diffusion.py index 8f04761502..5bf6674a43 100644 --- a/examples/community/llm_grounded_diffusion.py +++ b/examples/community/llm_grounded_diffusion.py @@ -302,7 +302,7 @@ class LLMGroundedDiffusionPipeline( [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. @@ -392,8 +392,8 @@ class LLMGroundedDiffusionPipeline( "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" + " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" diff --git a/examples/community/lpw_stable_diffusion.py b/examples/community/lpw_stable_diffusion.py index cb017c0bbe..58e932bbcf 100644 --- a/examples/community/lpw_stable_diffusion.py +++ b/examples/community/lpw_stable_diffusion.py @@ -552,8 +552,8 @@ class StableDiffusionLongPromptWeightingPipeline( "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" + " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" diff --git a/examples/community/lpw_stable_diffusion_xl.py b/examples/community/lpw_stable_diffusion_xl.py index 272c5d5652..95d3405df5 100644 --- a/examples/community/lpw_stable_diffusion_xl.py +++ b/examples/community/lpw_stable_diffusion_xl.py @@ -1765,7 +1765,7 @@ class SDXLLongPromptWeightingPipeline( # Check that sizes of mask, masked image and latents match if num_channels_unet == 9: - # default case for runwayml/stable-diffusion-inpainting + # default case for stable-diffusion-v1-5/stable-diffusion-inpainting num_channels_mask = mask.shape[1] num_channels_masked_image = masked_image_latents.shape[1] if num_channels_latents + num_channels_mask + num_channels_masked_image != num_channels_unet: diff --git a/examples/community/matryoshka.py b/examples/community/matryoshka.py index 3871552672..97ad8b9e86 100644 --- a/examples/community/matryoshka.py +++ b/examples/community/matryoshka.py @@ -3729,8 +3729,8 @@ class MatryoshkaPipeline( "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" + " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" diff --git a/examples/community/multilingual_stable_diffusion.py b/examples/community/multilingual_stable_diffusion.py index afef4e9e97..436803f201 100644 --- a/examples/community/multilingual_stable_diffusion.py +++ b/examples/community/multilingual_stable_diffusion.py @@ -78,7 +78,7 @@ class MultilingualStableDiffusion(DiffusionPipeline, StableDiffusionMixin): [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details. + Please, refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for details. feature_extractor ([`CLIPImageProcessor`]): Model that extracts features from generated images to be used as inputs for the `safety_checker`. """ diff --git a/examples/community/pipeline_controlnet_xl_kolors_inpaint.py b/examples/community/pipeline_controlnet_xl_kolors_inpaint.py index 4b6123cc1f..3abd984829 100644 --- a/examples/community/pipeline_controlnet_xl_kolors_inpaint.py +++ b/examples/community/pipeline_controlnet_xl_kolors_inpaint.py @@ -1607,7 +1607,7 @@ class KolorsControlNetInpaintPipeline( # 9. Check that sizes of mask, masked image and latents match if num_channels_unet == 9: - # default case for runwayml/stable-diffusion-inpainting + # default case for stable-diffusion-v1-5/stable-diffusion-inpainting num_channels_mask = mask.shape[1] num_channels_masked_image = masked_image_latents.shape[1] if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels: diff --git a/examples/community/pipeline_fabric.py b/examples/community/pipeline_fabric.py index dcc7730cbe..d29e98df5e 100644 --- a/examples/community/pipeline_fabric.py +++ b/examples/community/pipeline_fabric.py @@ -135,7 +135,7 @@ class FabricPipeline(DiffusionPipeline): [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for more details about a model's potential harms. """ @@ -163,8 +163,8 @@ class FabricPipeline(DiffusionPipeline): "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" + " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" diff --git a/examples/community/pipeline_kolors_inpainting.py b/examples/community/pipeline_kolors_inpainting.py index 3cab8ecac0..26517819eb 100644 --- a/examples/community/pipeline_kolors_inpainting.py +++ b/examples/community/pipeline_kolors_inpainting.py @@ -1487,7 +1487,7 @@ class KolorsInpaintPipeline( # 8. Check that sizes of mask, masked image and latents match if num_channels_unet == 9: - # default case for runwayml/stable-diffusion-inpainting + # default case for stable-diffusion-v1-5/stable-diffusion-inpainting num_channels_mask = mask.shape[1] num_channels_masked_image = masked_image_latents.shape[1] if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels: diff --git a/examples/community/pipeline_prompt2prompt.py b/examples/community/pipeline_prompt2prompt.py index 8d94dc9248..eb19667970 100644 --- a/examples/community/pipeline_prompt2prompt.py +++ b/examples/community/pipeline_prompt2prompt.py @@ -106,7 +106,7 @@ class Prompt2PromptPipeline( [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. @@ -187,8 +187,8 @@ class Prompt2PromptPipeline( "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" + " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" diff --git a/examples/community/pipeline_sdxl_style_aligned.py b/examples/community/pipeline_sdxl_style_aligned.py index 10438af365..51547599f5 100644 --- a/examples/community/pipeline_sdxl_style_aligned.py +++ b/examples/community/pipeline_sdxl_style_aligned.py @@ -1730,7 +1730,7 @@ class StyleAlignedSDXLPipeline( # Check that sizes of mask, masked image and latents match if num_channels_unet == 9: - # default case for runwayml/stable-diffusion-inpainting + # default case for stable-diffusion-v1-5/stable-diffusion-inpainting num_channels_mask = mask.shape[1] num_channels_masked_image = masked_image_latents.shape[1] if num_channels_latents + num_channels_mask + num_channels_masked_image != num_channels_unet: diff --git a/examples/community/pipeline_stable_diffusion_boxdiff.py b/examples/community/pipeline_stable_diffusion_boxdiff.py index 07e29b9c05..c05a614313 100644 --- a/examples/community/pipeline_stable_diffusion_boxdiff.py +++ b/examples/community/pipeline_stable_diffusion_boxdiff.py @@ -59,7 +59,7 @@ EXAMPLE_DOC_STRING = """ >>> import torch >>> from diffusers import StableDiffusionPipeline - >>> pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16) + >>> pipe = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16) >>> pipe = pipe.to("cuda") >>> prompt = "a photo of an astronaut riding a horse on mars" @@ -392,7 +392,7 @@ class StableDiffusionBoxDiffPipeline( [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. @@ -473,8 +473,8 @@ class StableDiffusionBoxDiffPipeline( "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" + " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" diff --git a/examples/community/pipeline_stable_diffusion_pag.py b/examples/community/pipeline_stable_diffusion_pag.py index 6b62b610af..3f98dca0b9 100644 --- a/examples/community/pipeline_stable_diffusion_pag.py +++ b/examples/community/pipeline_stable_diffusion_pag.py @@ -42,7 +42,7 @@ EXAMPLE_DOC_STRING = """ ```py >>> import torch >>> from diffusers import StableDiffusionPipeline - >>> pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16) + >>> pipe = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16) >>> pipe = pipe.to("cuda") >>> prompt = "a photo of an astronaut riding a horse on mars" >>> image = pipe(prompt).images[0] @@ -359,7 +359,7 @@ class StableDiffusionPAGPipeline( [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. @@ -440,8 +440,8 @@ class StableDiffusionPAGPipeline( "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" + " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" diff --git a/examples/community/pipeline_stable_diffusion_upscale_ldm3d.py b/examples/community/pipeline_stable_diffusion_upscale_ldm3d.py index 9777633535..e358f66b4a 100644 --- a/examples/community/pipeline_stable_diffusion_upscale_ldm3d.py +++ b/examples/community/pipeline_stable_diffusion_upscale_ldm3d.py @@ -100,7 +100,7 @@ class StableDiffusionUpscaleLDM3DPipeline( [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. diff --git a/examples/community/pipeline_stable_diffusion_xl_attentive_eraser.py b/examples/community/pipeline_stable_diffusion_xl_attentive_eraser.py index a881814c2a..65df4c03eb 100644 --- a/examples/community/pipeline_stable_diffusion_xl_attentive_eraser.py +++ b/examples/community/pipeline_stable_diffusion_xl_attentive_eraser.py @@ -2042,7 +2042,7 @@ class StableDiffusionXL_AE_Pipeline( # 8. Check that sizes of mask, masked image and latents match if num_channels_unet == 9: - # default case for runwayml/stable-diffusion-inpainting + # default case for stable-diffusion-v1-5/stable-diffusion-inpainting num_channels_mask = mask.shape[1] num_channels_masked_image = masked_image_latents.shape[1] if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels: diff --git a/examples/community/pipeline_stable_diffusion_xl_controlnet_adapter.py b/examples/community/pipeline_stable_diffusion_xl_controlnet_adapter.py index 564a19e923..6dade126f2 100644 --- a/examples/community/pipeline_stable_diffusion_xl_controlnet_adapter.py +++ b/examples/community/pipeline_stable_diffusion_xl_controlnet_adapter.py @@ -188,7 +188,7 @@ class StableDiffusionXLControlNetAdapterPipeline( [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details. + Please, refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for details. feature_extractor ([`CLIPImageProcessor`]): Model that extracts features from generated images to be used as inputs for the `safety_checker`. """ diff --git a/examples/community/pipeline_stable_diffusion_xl_controlnet_adapter_inpaint.py b/examples/community/pipeline_stable_diffusion_xl_controlnet_adapter_inpaint.py index c73433b20f..9ec6a90b07 100644 --- a/examples/community/pipeline_stable_diffusion_xl_controlnet_adapter_inpaint.py +++ b/examples/community/pipeline_stable_diffusion_xl_controlnet_adapter_inpaint.py @@ -330,7 +330,7 @@ class StableDiffusionXLControlNetAdapterInpaintPipeline( [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details. + Please, refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for details. feature_extractor ([`CLIPImageProcessor`]): Model that extracts features from generated images to be used as inputs for the `safety_checker`. requires_aesthetics_score (`bool`, *optional*, defaults to `"False"`): @@ -1569,7 +1569,7 @@ class StableDiffusionXLControlNetAdapterInpaintPipeline( # 8. Check that sizes of mask, masked image and latents match if num_channels_unet == 9: - # default case for runwayml/stable-diffusion-inpainting + # default case for stable-diffusion-v1-5/stable-diffusion-inpainting num_channels_mask = mask.shape[1] num_channels_masked_image = masked_image_latents.shape[1] if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels: diff --git a/examples/community/pipeline_zero1to3.py b/examples/community/pipeline_zero1to3.py index 9e29566978..1be59fd832 100644 --- a/examples/community/pipeline_zero1to3.py +++ b/examples/community/pipeline_zero1to3.py @@ -46,7 +46,7 @@ EXAMPLE_DOC_STRING = """ >>> import torch >>> from diffusers import StableDiffusionPipeline - >>> pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16) + >>> pipe = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16) >>> pipe = pipe.to("cuda") >>> prompt = "a photo of an astronaut riding a horse on mars" @@ -86,7 +86,7 @@ class Zero1to3StableDiffusionPipeline(DiffusionPipeline, StableDiffusionMixin): [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details. + Please, refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for details. feature_extractor ([`CLIPImageProcessor`]): Model that extracts features from generated images to be used as inputs for the `safety_checker`. cc_projection ([`CCProjection`]): @@ -164,8 +164,8 @@ class Zero1to3StableDiffusionPipeline(DiffusionPipeline, StableDiffusionMixin): "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" + " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" diff --git a/examples/community/rerender_a_video.py b/examples/community/rerender_a_video.py index 78a15a03b0..840f9e206d 100644 --- a/examples/community/rerender_a_video.py +++ b/examples/community/rerender_a_video.py @@ -288,7 +288,7 @@ class RerenderAVideoPipeline(StableDiffusionControlNetImg2ImgPipeline): [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details. + Please, refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for details. feature_extractor ([`CLIPImageProcessor`]): Model that extracts features from generated images to be used as inputs for the `safety_checker`. """ diff --git a/examples/community/run_onnx_controlnet.py b/examples/community/run_onnx_controlnet.py index f0ab2a2b96..2b56e8a1e5 100644 --- a/examples/community/run_onnx_controlnet.py +++ b/examples/community/run_onnx_controlnet.py @@ -54,7 +54,7 @@ EXAMPLE_DOC_STRING = """ >>> # load control net and stable diffusion v1-5 >>> controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16) >>> pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained( - ... "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16 + ... "stable-diffusion-v1-5/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16 ... ) >>> # speed up diffusion process with faster scheduler and memory optimization diff --git a/examples/community/run_tensorrt_controlnet.py b/examples/community/run_tensorrt_controlnet.py index e4f1abc83b..b62eb4f58e 100644 --- a/examples/community/run_tensorrt_controlnet.py +++ b/examples/community/run_tensorrt_controlnet.py @@ -158,7 +158,7 @@ EXAMPLE_DOC_STRING = """ >>> # load control net and stable diffusion v1-5 >>> controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16) >>> pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained( - ... "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16 + ... "stable-diffusion-v1-5/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16 ... ) >>> # speed up diffusion process with faster scheduler and memory optimization diff --git a/examples/community/sd_text2img_k_diffusion.py b/examples/community/sd_text2img_k_diffusion.py index 4d5cea497f..e351420f78 100755 --- a/examples/community/sd_text2img_k_diffusion.py +++ b/examples/community/sd_text2img_k_diffusion.py @@ -64,7 +64,7 @@ class StableDiffusionPipeline(DiffusionPipeline, StableDiffusionMixin): [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details. + Please, refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for details. feature_extractor ([`CLIPImageProcessor`]): Model that extracts features from generated images to be used as inputs for the `safety_checker`. """ diff --git a/examples/community/sde_drag.py b/examples/community/sde_drag.py index f408ee64db..63899ce738 100644 --- a/examples/community/sde_drag.py +++ b/examples/community/sde_drag.py @@ -114,7 +114,7 @@ class SdeDragPipeline(DiffusionPipeline): >>> from diffusers import DDIMScheduler, DiffusionPipeline >>> # Load the pipeline - >>> model_path = "runwayml/stable-diffusion-v1-5" + >>> model_path = "stable-diffusion-v1-5/stable-diffusion-v1-5" >>> scheduler = DDIMScheduler.from_pretrained(model_path, subfolder="scheduler") >>> pipe = DiffusionPipeline.from_pretrained(model_path, scheduler=scheduler, custom_pipeline="sde_drag") >>> pipe.to('cuda') diff --git a/examples/community/stable_diffusion_comparison.py b/examples/community/stable_diffusion_comparison.py index 22f3b3e0c3..ce6e77c87f 100644 --- a/examples/community/stable_diffusion_comparison.py +++ b/examples/community/stable_diffusion_comparison.py @@ -46,7 +46,7 @@ class StableDiffusionComparisonPipeline(DiffusionPipeline, StableDiffusionMixin) [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionMegaSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details. + Please, refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for details. feature_extractor ([`CLIPImageProcessor`]): Model that extracts features from generated images to be used as inputs for the `safety_checker`. """ diff --git a/examples/community/stable_diffusion_controlnet_img2img.py b/examples/community/stable_diffusion_controlnet_img2img.py index 6d8038cfd4..aa116112be 100644 --- a/examples/community/stable_diffusion_controlnet_img2img.py +++ b/examples/community/stable_diffusion_controlnet_img2img.py @@ -36,7 +36,7 @@ EXAMPLE_DOC_STRING = """ >>> controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16) >>> pipe_controlnet = StableDiffusionControlNetImg2ImgPipeline.from_pretrained( - "runwayml/stable-diffusion-v1-5", + "stable-diffusion-v1-5/stable-diffusion-v1-5", controlnet=controlnet, safety_checker=None, torch_dtype=torch.float16 diff --git a/examples/community/stable_diffusion_controlnet_inpaint.py b/examples/community/stable_diffusion_controlnet_inpaint.py index fe7b808b6b..6d710e0d73 100644 --- a/examples/community/stable_diffusion_controlnet_inpaint.py +++ b/examples/community/stable_diffusion_controlnet_inpaint.py @@ -81,7 +81,7 @@ EXAMPLE_DOC_STRING = """ >>> controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-seg", torch_dtype=torch.float16) >>> pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained( - "runwayml/stable-diffusion-inpainting", controlnet=controlnet, safety_checker=None, torch_dtype=torch.float16 + "stable-diffusion-v1-5/stable-diffusion-inpainting", controlnet=controlnet, safety_checker=None, torch_dtype=torch.float16 ) >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) diff --git a/examples/community/stable_diffusion_controlnet_inpaint_img2img.py b/examples/community/stable_diffusion_controlnet_inpaint_img2img.py index 2b5dc77fe5..fcb5ed059b 100644 --- a/examples/community/stable_diffusion_controlnet_inpaint_img2img.py +++ b/examples/community/stable_diffusion_controlnet_inpaint_img2img.py @@ -80,7 +80,7 @@ EXAMPLE_DOC_STRING = """ >>> controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-seg", torch_dtype=torch.float16) >>> pipe = StableDiffusionControlNetInpaintImg2ImgPipeline.from_pretrained( - "runwayml/stable-diffusion-inpainting", controlnet=controlnet, safety_checker=None, torch_dtype=torch.float16 + "stable-diffusion-v1-5/stable-diffusion-inpainting", controlnet=controlnet, safety_checker=None, torch_dtype=torch.float16 ) >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) diff --git a/examples/community/stable_diffusion_controlnet_reference.py b/examples/community/stable_diffusion_controlnet_reference.py index e5dd249e04..74c81b6362 100644 --- a/examples/community/stable_diffusion_controlnet_reference.py +++ b/examples/community/stable_diffusion_controlnet_reference.py @@ -37,7 +37,7 @@ EXAMPLE_DOC_STRING = """ >>> controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16) >>> pipe = StableDiffusionControlNetReferencePipeline.from_pretrained( - "runwayml/stable-diffusion-v1-5", + "stable-diffusion-v1-5/stable-diffusion-v1-5", controlnet=controlnet, safety_checker=None, torch_dtype=torch.float16 diff --git a/examples/community/stable_diffusion_ipex.py b/examples/community/stable_diffusion_ipex.py index 7d1cd4f5d0..4f545aa09d 100644 --- a/examples/community/stable_diffusion_ipex.py +++ b/examples/community/stable_diffusion_ipex.py @@ -43,7 +43,7 @@ EXAMPLE_DOC_STRING = """ >>> import torch >>> from diffusers import StableDiffusionPipeline - >>> pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", custom_pipeline="stable_diffusion_ipex") + >>> pipe = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", custom_pipeline="stable_diffusion_ipex") >>> # For Float32 >>> pipe.prepare_for_ipex(prompt, dtype=torch.float32, height=512, width=512) #value of image height/width should be consistent with the pipeline inference @@ -85,7 +85,7 @@ class StableDiffusionIPEXPipeline( [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details. + Please, refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for details. feature_extractor ([`CLIPImageProcessor`]): Model that extracts features from generated images to be used as inputs for the `safety_checker`. """ @@ -161,8 +161,8 @@ class StableDiffusionIPEXPipeline( "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" + " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" diff --git a/examples/community/stable_diffusion_mega.py b/examples/community/stable_diffusion_mega.py index 77e5011d2a..c67ebc80b0 100644 --- a/examples/community/stable_diffusion_mega.py +++ b/examples/community/stable_diffusion_mega.py @@ -47,7 +47,7 @@ class StableDiffusionMegaPipeline(DiffusionPipeline, StableDiffusionMixin): [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionMegaSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details. + Please, refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for details. feature_extractor ([`CLIPImageProcessor`]): Model that extracts features from generated images to be used as inputs for the `safety_checker`. """ diff --git a/examples/community/stable_diffusion_reference.py b/examples/community/stable_diffusion_reference.py index 6f7dce9823..d0372bbeba 100644 --- a/examples/community/stable_diffusion_reference.py +++ b/examples/community/stable_diffusion_reference.py @@ -46,7 +46,7 @@ EXAMPLE_DOC_STRING = """ >>> input_image = load_image("https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png") >>> pipe = StableDiffusionReferencePipeline.from_pretrained( - "runwayml/stable-diffusion-v1-5", + "stable-diffusion-v1-5/stable-diffusion-v1-5", safety_checker=None, torch_dtype=torch.float16 ).to('cuda:0') @@ -112,7 +112,7 @@ class StableDiffusionReferencePipeline( [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details. + Please, refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for details. feature_extractor ([`CLIPImageProcessor`]): Model that extracts features from generated images to be used as inputs for the `safety_checker`. """ @@ -194,8 +194,8 @@ class StableDiffusionReferencePipeline( "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely .If you're checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" + " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" diff --git a/examples/community/stable_diffusion_repaint.py b/examples/community/stable_diffusion_repaint.py index 94b9f8b01b..b974e3c7ae 100644 --- a/examples/community/stable_diffusion_repaint.py +++ b/examples/community/stable_diffusion_repaint.py @@ -167,7 +167,7 @@ class StableDiffusionRepaintPipeline( [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details. + Please, refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for details. feature_extractor ([`CLIPImageProcessor`]): Model that extracts features from generated images to be used as inputs for the `safety_checker`. """ @@ -249,8 +249,8 @@ class StableDiffusionRepaintPipeline( "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely .If you're checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" + " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" diff --git a/examples/community/stable_diffusion_tensorrt_img2img.py b/examples/community/stable_diffusion_tensorrt_img2img.py index dc11703b6a..5b7733fe57 100755 --- a/examples/community/stable_diffusion_tensorrt_img2img.py +++ b/examples/community/stable_diffusion_tensorrt_img2img.py @@ -678,7 +678,7 @@ class TensorRTStableDiffusionImg2ImgPipeline(DiffusionPipeline): [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details. + Please, refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for details. feature_extractor ([`CLIPImageProcessor`]): Model that extracts features from generated images to be used as inputs for the `safety_checker`. """ @@ -766,8 +766,8 @@ class TensorRTStableDiffusionImg2ImgPipeline(DiffusionPipeline): "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" + " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" diff --git a/examples/community/stable_diffusion_tensorrt_inpaint.py b/examples/community/stable_diffusion_tensorrt_inpaint.py index fff7309e9c..fc81e4c289 100755 --- a/examples/community/stable_diffusion_tensorrt_inpaint.py +++ b/examples/community/stable_diffusion_tensorrt_inpaint.py @@ -682,7 +682,7 @@ class TensorRTStableDiffusionInpaintPipeline(DiffusionPipeline): [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details. + Please, refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for details. feature_extractor ([`CLIPImageProcessor`]): Model that extracts features from generated images to be used as inputs for the `safety_checker`. """ @@ -770,8 +770,8 @@ class TensorRTStableDiffusionInpaintPipeline(DiffusionPipeline): "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" + " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" diff --git a/examples/community/stable_diffusion_tensorrt_txt2img.py b/examples/community/stable_diffusion_tensorrt_txt2img.py index 15a6e69c41..e1d09edf93 100755 --- a/examples/community/stable_diffusion_tensorrt_txt2img.py +++ b/examples/community/stable_diffusion_tensorrt_txt2img.py @@ -594,7 +594,7 @@ class TensorRTStableDiffusionPipeline(DiffusionPipeline): [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details. + Please, refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for details. feature_extractor ([`CLIPImageProcessor`]): Model that extracts features from generated images to be used as inputs for the `safety_checker`. """ @@ -682,8 +682,8 @@ class TensorRTStableDiffusionPipeline(DiffusionPipeline): "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" + " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" diff --git a/examples/community/text_inpainting.py b/examples/community/text_inpainting.py index f262cf2cac..bdf9eca498 100644 --- a/examples/community/text_inpainting.py +++ b/examples/community/text_inpainting.py @@ -52,7 +52,7 @@ class TextInpainting(DiffusionPipeline, StableDiffusionMixin): [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details. + Please, refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for details. feature_extractor ([`CLIPImageProcessor`]): Model that extracts features from generated images to be used as inputs for the `safety_checker`. """ diff --git a/examples/research_projects/anytext/anytext.py b/examples/research_projects/anytext/anytext.py index 38f0adb891..7ae6ae57c2 100644 --- a/examples/research_projects/anytext/anytext.py +++ b/examples/research_projects/anytext/anytext.py @@ -1223,7 +1223,7 @@ class AnyTextPipeline( [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. diff --git a/examples/research_projects/dreambooth_inpaint/README.md b/examples/research_projects/dreambooth_inpaint/README.md index 46703fa982..b6ee1d72f6 100644 --- a/examples/research_projects/dreambooth_inpaint/README.md +++ b/examples/research_projects/dreambooth_inpaint/README.md @@ -5,7 +5,7 @@ This script was added by @thedarkzeno . Please note that this script is not actively maintained, you can open an issue and tag @thedarkzeno or @patil-suraj though. ```bash -export MODEL_NAME="runwayml/stable-diffusion-inpainting" +export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-inpainting" export INSTANCE_DIR="path-to-instance-images" export OUTPUT_DIR="path-to-save-model" @@ -29,7 +29,7 @@ Prior-preservation is used to avoid overfitting and language-drift. Refer to the According to the paper, it's recommended to generate `num_epochs * num_samples` images for prior-preservation. 200-300 works well for most cases. ```bash -export MODEL_NAME="runwayml/stable-diffusion-inpainting" +export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-inpainting" export INSTANCE_DIR="path-to-instance-images" export CLASS_DIR="path-to-class-images" export OUTPUT_DIR="path-to-save-model" @@ -60,7 +60,7 @@ With the help of gradient checkpointing and the 8-bit optimizer from bitsandbyte To install `bitandbytes` please refer to this [readme](https://github.com/TimDettmers/bitsandbytes#requirements--installation). ```bash -export MODEL_NAME="runwayml/stable-diffusion-inpainting" +export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-inpainting" export INSTANCE_DIR="path-to-instance-images" export CLASS_DIR="path-to-class-images" export OUTPUT_DIR="path-to-save-model" @@ -92,7 +92,7 @@ Pass the `--train_text_encoder` argument to the script to enable training `text_ ___Note: Training text encoder requires more memory, with this option the training won't fit on 16GB GPU. It needs at least 24GB VRAM.___ ```bash -export MODEL_NAME="runwayml/stable-diffusion-inpainting" +export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-inpainting" export INSTANCE_DIR="path-to-instance-images" export CLASS_DIR="path-to-class-images" export OUTPUT_DIR="path-to-save-model" diff --git a/examples/research_projects/ip_adapter/README.md b/examples/research_projects/ip_adapter/README.md index 3df9644ddf..0bead5ae85 100644 --- a/examples/research_projects/ip_adapter/README.md +++ b/examples/research_projects/ip_adapter/README.md @@ -55,7 +55,7 @@ The Accelerate launch command is used to train a model using multiple GPUs and m ``` accelerate launch --mixed_precision "fp16" \ tutorial_train_ip-adapter.py \ ---pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5/" \ +--pretrained_model_name_or_path="stable-diffusion-v1-5/stable-diffusion-v1-5/" \ --image_encoder_path="{image_encoder_path}" \ --data_json_file="{data.json}" \ --data_root_path="{image_path}" \ @@ -73,7 +73,7 @@ tutorial_train_ip-adapter.py \ ``` accelerate launch --num_processes 8 --multi_gpu --mixed_precision "fp16" \ tutorial_train_ip-adapter.py \ - --pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5/" \ + --pretrained_model_name_or_path="stable-diffusion-v1-5/stable-diffusion-v1-5/" \ --image_encoder_path="{image_encoder_path}" \ --data_json_file="{data.json}" \ --data_root_path="{image_path}" \ diff --git a/examples/research_projects/multi_subject_dreambooth_inpainting/README.md b/examples/research_projects/multi_subject_dreambooth_inpainting/README.md index 8ddef1b83c..3412de662f 100644 --- a/examples/research_projects/multi_subject_dreambooth_inpainting/README.md +++ b/examples/research_projects/multi_subject_dreambooth_inpainting/README.md @@ -27,7 +27,7 @@ You can build multiple datasets for every subject and upload them to the 🤗 hu Before launching the training script, make sure to select the inpainting the target model, the output directory and the 🤗 datasets. ```bash -export MODEL_NAME="runwayml/stable-diffusion-inpainting" +export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-inpainting" export OUTPUT_DIR="path-to-save-model" export DATASET_1="gzguevara/mr_potato_head_masked" diff --git a/examples/research_projects/promptdiffusion/pipeline_prompt_diffusion.py b/examples/research_projects/promptdiffusion/pipeline_prompt_diffusion.py index 1bd9c0161f..233df12765 100644 --- a/examples/research_projects/promptdiffusion/pipeline_prompt_diffusion.py +++ b/examples/research_projects/promptdiffusion/pipeline_prompt_diffusion.py @@ -177,7 +177,7 @@ class PromptDiffusionPipeline( [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. diff --git a/examples/research_projects/vae/vae_roundtrip.py b/examples/research_projects/vae/vae_roundtrip.py index cdc3a54fdf..922cb42615 100644 --- a/examples/research_projects/vae/vae_roundtrip.py +++ b/examples/research_projects/vae/vae_roundtrip.py @@ -238,7 +238,7 @@ def parse_args() -> argparse.Namespace: # EXAMPLE USAGE: # -# python vae_roundtrip.py --use_cuda --pretrained_model_name_or_path "runwayml/stable-diffusion-v1-5" --subfolder "vae" --input_image "foo.png" +# python vae_roundtrip.py --use_cuda --pretrained_model_name_or_path "stable-diffusion-v1-5/stable-diffusion-v1-5" --subfolder "vae" --input_image "foo.png" # # python vae_roundtrip.py --use_cuda --pretrained_model_name_or_path "madebyollin/taesd" --use_tiny_nn --input_image "foo.png" # diff --git a/scripts/convert_consistency_decoder.py b/scripts/convert_consistency_decoder.py index 629c784c09..9e28945775 100644 --- a/scripts/convert_consistency_decoder.py +++ b/scripts/convert_consistency_decoder.py @@ -24,7 +24,8 @@ args = args.parse_args() def _extract_into_tensor(arr, timesteps, broadcast_shape): - # from: https://github.com/openai/guided-diffusion/blob/22e0df8183507e13a7813f8d38d51b072ca1e67c/guided_diffusion/gaussian_diffusion.py#L895 """ + # from: https://github.com/openai/guided-diffusion/blob/22e0df8183507e13a7813f8d38d51b072ca1e67c/guided_diffusion/gaussian_diffusion.py#L895 + # """ res = arr[timesteps].float() dims_to_append = len(broadcast_shape) - len(res.shape) return res[(...,) + (None,) * dims_to_append] @@ -507,7 +508,9 @@ def rename_state_dict(sd, embedding): # encode with stable diffusion vae -pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16) +pipe = StableDiffusionPipeline.from_pretrained( + "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16 +) pipe.vae.cuda() # construct original decoder with jitted model @@ -1090,7 +1093,7 @@ def new_constructor(self, **kwargs): Encoder.__init__ = new_constructor -vae = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae") +vae = AutoencoderKL.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="vae") consistency_vae = ConsistencyDecoderVAE( encoder_args=vae.encoder.constructor_arguments, decoder_args=unet.config, @@ -1117,7 +1120,7 @@ print((sample_consistency_orig - sample_consistency_new_3).abs().sum()) print("running with diffusers pipeline") pipe = DiffusionPipeline.from_pretrained( - "runwayml/stable-diffusion-v1-5", vae=consistency_vae, torch_dtype=torch.float16 + "stable-diffusion-v1-5/stable-diffusion-v1-5", vae=consistency_vae, torch_dtype=torch.float16 ) pipe.to("cuda") diff --git a/src/diffusers/models/auto_model.py b/src/diffusers/models/auto_model.py index a95b0ae64a..947b610ea6 100644 --- a/src/diffusers/models/auto_model.py +++ b/src/diffusers/models/auto_model.py @@ -128,13 +128,13 @@ class AutoModel(ConfigMixin): ```py from diffusers import AutoModel - unet = AutoModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet") + unet = AutoModel.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="unet") ``` If you get the error message below, you need to finetune the weights for your downstream task: ```bash - Some weights of UNet2DConditionModel were not initialized from the model checkpoint at runwayml/stable-diffusion-v1-5 and are newly initialized because the shapes did not match: + Some weights of UNet2DConditionModel were not initialized from the model checkpoint at stable-diffusion-v1-5/stable-diffusion-v1-5 and are newly initialized because the shapes did not match: - conv_in.weight: found shape torch.Size([320, 4, 3, 3]) in the checkpoint and torch.Size([320, 9, 3, 3]) in the model instantiated You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference. ``` diff --git a/src/diffusers/models/modeling_flax_utils.py b/src/diffusers/models/modeling_flax_utils.py index fd19578321..3f06099319 100644 --- a/src/diffusers/models/modeling_flax_utils.py +++ b/src/diffusers/models/modeling_flax_utils.py @@ -113,14 +113,14 @@ class FlaxModelMixin(PushToHubMixin): >>> from diffusers import FlaxUNet2DConditionModel >>> # load model - >>> model, params = FlaxUNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5") + >>> model, params = FlaxUNet2DConditionModel.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5") >>> # By default, the model parameters will be in fp32 precision, to cast these to bfloat16 precision >>> params = model.to_bf16(params) >>> # If you don't want to cast certain parameters (for example layer norm bias and scale) >>> # then pass the mask as follows >>> from flax import traverse_util - >>> model, params = FlaxUNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5") + >>> model, params = FlaxUNet2DConditionModel.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5") >>> flat_params = traverse_util.flatten_dict(params) >>> mask = { ... path: (path[-2] != ("LayerNorm", "bias") and path[-2:] != ("LayerNorm", "scale")) @@ -149,7 +149,7 @@ class FlaxModelMixin(PushToHubMixin): >>> from diffusers import FlaxUNet2DConditionModel >>> # Download model and configuration from huggingface.co - >>> model, params = FlaxUNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5") + >>> model, params = FlaxUNet2DConditionModel.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5") >>> # By default, the model params will be in fp32, to illustrate the use of this method, >>> # we'll first cast to fp16 and back to fp32 >>> params = model.to_f16(params) @@ -179,14 +179,14 @@ class FlaxModelMixin(PushToHubMixin): >>> from diffusers import FlaxUNet2DConditionModel >>> # load model - >>> model, params = FlaxUNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5") + >>> model, params = FlaxUNet2DConditionModel.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5") >>> # By default, the model params will be in fp32, to cast these to float16 >>> params = model.to_fp16(params) >>> # If you want don't want to cast certain parameters (for example layer norm bias and scale) >>> # then pass the mask as follows >>> from flax import traverse_util - >>> model, params = FlaxUNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5") + >>> model, params = FlaxUNet2DConditionModel.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5") >>> flat_params = traverse_util.flatten_dict(params) >>> mask = { ... path: (path[-2] != ("LayerNorm", "bias") and path[-2:] != ("LayerNorm", "scale")) @@ -216,8 +216,8 @@ class FlaxModelMixin(PushToHubMixin): pretrained_model_name_or_path (`str` or `os.PathLike`): Can be either: - - A string, the *model id* (for example `runwayml/stable-diffusion-v1-5`) of a pretrained model - hosted on the Hub. + - A string, the *model id* (for example `stable-diffusion-v1-5/stable-diffusion-v1-5`) of a + pretrained model hosted on the Hub. - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved using [`~FlaxModelMixin.save_pretrained`]. dtype (`jax.numpy.dtype`, *optional*, defaults to `jax.numpy.float32`): @@ -271,7 +271,7 @@ class FlaxModelMixin(PushToHubMixin): >>> from diffusers import FlaxUNet2DConditionModel >>> # Download model and configuration from huggingface.co and cache. - >>> model, params = FlaxUNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5") + >>> model, params = FlaxUNet2DConditionModel.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5") >>> # Model was saved using *save_pretrained('./test/saved_model/')* (for example purposes, not runnable). >>> model, params = FlaxUNet2DConditionModel.from_pretrained("./test/saved_model/") ``` @@ -279,7 +279,7 @@ class FlaxModelMixin(PushToHubMixin): If you get the error message below, you need to finetune the weights for your downstream task: ```bash - Some weights of UNet2DConditionModel were not initialized from the model checkpoint at runwayml/stable-diffusion-v1-5 and are newly initialized because the shapes did not match: + Some weights of UNet2DConditionModel were not initialized from the model checkpoint at stable-diffusion-v1-5/stable-diffusion-v1-5 and are newly initialized because the shapes did not match: - conv_in.weight: found shape torch.Size([320, 4, 3, 3]) in the checkpoint and torch.Size([320, 9, 3, 3]) in the model instantiated You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference. ``` diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index 1af7ba9ac5..ce26785f63 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -923,13 +923,13 @@ class ModelMixin(torch.nn.Module, PushToHubMixin): ```py from diffusers import UNet2DConditionModel - unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet") + unet = UNet2DConditionModel.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="unet") ``` If you get the error message below, you need to finetune the weights for your downstream task: ```bash - Some weights of UNet2DConditionModel were not initialized from the model checkpoint at runwayml/stable-diffusion-v1-5 and are newly initialized because the shapes did not match: + Some weights of UNet2DConditionModel were not initialized from the model checkpoint at stable-diffusion-v1-5/stable-diffusion-v1-5 and are newly initialized because the shapes did not match: - conv_in.weight: found shape torch.Size([320, 4, 3, 3]) in the checkpoint and torch.Size([320, 9, 3, 3]) in the model instantiated You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference. ``` @@ -1800,7 +1800,7 @@ class ModelMixin(torch.nn.Module, PushToHubMixin): ```py from diffusers import UNet2DConditionModel - model_id = "runwayml/stable-diffusion-v1-5" + model_id = "stable-diffusion-v1-5/stable-diffusion-v1-5" unet = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet") unet.num_parameters(only_trainable=True) 859520964 diff --git a/src/diffusers/modular_pipelines/stable_diffusion_xl/denoise.py b/src/diffusers/modular_pipelines/stable_diffusion_xl/denoise.py index 8a80257473..83e6007e3f 100644 --- a/src/diffusers/modular_pipelines/stable_diffusion_xl/denoise.py +++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/denoise.py @@ -115,7 +115,7 @@ class StableDiffusionXLInpaintLoopBeforeDenoiser(ModularPipelineBlocks): def check_inputs(components, block_state): num_channels_unet = components.num_channels_unet if num_channels_unet == 9: - # default case for runwayml/stable-diffusion-inpainting + # default case for stable-diffusion-v1-5/stable-diffusion-inpainting if block_state.mask is None or block_state.masked_image_latents is None: raise ValueError("mask and masked_image_latents must be provided for inpainting-specific Unet") num_channels_latents = block_state.latents.shape[1] diff --git a/src/diffusers/pipelines/README.md b/src/diffusers/pipelines/README.md index 363caffe20..6f9ab7b291 100644 --- a/src/diffusers/pipelines/README.md +++ b/src/diffusers/pipelines/README.md @@ -159,7 +159,7 @@ init_image = download_image(img_url).resize((512, 512)) mask_image = download_image(mask_url).resize((512, 512)) pipe = StableDiffusionInpaintPipeline.from_pretrained( - "runwayml/stable-diffusion-inpainting", + "stable-diffusion-v1-5/stable-diffusion-inpainting", torch_dtype=torch.float16, ) pipe = pipe.to("cuda") diff --git a/src/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs.py b/src/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs.py index 1545027a28..3682ddc911 100644 --- a/src/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs.py +++ b/src/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs.py @@ -133,8 +133,8 @@ class StableDiffusionControlNetXSPipeline( [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for + more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. """ diff --git a/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion.py b/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion.py index 003e748274..6f484aa3e2 100644 --- a/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion.py +++ b/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion.py @@ -185,8 +185,8 @@ class AltDiffusionPipeline( [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for + more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. """ @@ -266,8 +266,8 @@ class AltDiffusionPipeline( "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" + " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" diff --git a/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion_img2img.py b/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion_img2img.py index 64bd06d02e..d6bf901207 100644 --- a/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion_img2img.py +++ b/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion_img2img.py @@ -213,8 +213,8 @@ class AltDiffusionImg2ImgPipeline( [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for + more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. """ @@ -294,8 +294,8 @@ class AltDiffusionImg2ImgPipeline( "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" + " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" diff --git a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_cycle_diffusion.py b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_cycle_diffusion.py index 59c79e134e..08f8c7e26f 100644 --- a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_cycle_diffusion.py +++ b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_cycle_diffusion.py @@ -162,8 +162,8 @@ class CycleDiffusionPipeline(DiffusionPipeline, TextualInversionLoaderMixin, Sta instance of [`DDIMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for + more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. """ @@ -226,8 +226,8 @@ class CycleDiffusionPipeline(DiffusionPipeline, TextualInversionLoaderMixin, Sta "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely .If you're checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" + " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" diff --git a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_onnx_stable_diffusion_inpaint_legacy.py b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_onnx_stable_diffusion_inpaint_legacy.py index 2d9eaa493f..fcd8bf317a 100644 --- a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_onnx_stable_diffusion_inpaint_legacy.py +++ b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_onnx_stable_diffusion_inpaint_legacy.py @@ -62,7 +62,8 @@ class OnnxStableDiffusionInpaintPipelineLegacy(DiffusionPipeline): [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details. + Please, refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for + details. feature_extractor ([`CLIPImageProcessor`]): Model that extracts features from generated images to be used as inputs for the `safety_checker`. """ diff --git a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_inpaint_legacy.py b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_inpaint_legacy.py index 205ace65ee..ba0dd66c29 100644 --- a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_inpaint_legacy.py +++ b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_inpaint_legacy.py @@ -111,7 +111,8 @@ class StableDiffusionInpaintPipelineLegacy( [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details. + Please, refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for + details. feature_extractor ([`CLIPImageProcessor`]): Model that extracts features from generated images to be used as inputs for the `safety_checker`. """ @@ -196,8 +197,8 @@ class StableDiffusionInpaintPipelineLegacy( "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" + " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" diff --git a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_model_editing.py b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_model_editing.py index d81f0d2625..b7a0be57c1 100644 --- a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_model_editing.py +++ b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_model_editing.py @@ -64,8 +64,8 @@ class StableDiffusionModelEditingPipeline( [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for + more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. with_to_k ([`bool`]): diff --git a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_paradigms.py b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_paradigms.py index 2331157ba5..c236e73bf4 100644 --- a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_paradigms.py +++ b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_paradigms.py @@ -46,10 +46,12 @@ EXAMPLE_DOC_STRING = """ >>> from diffusers import DDPMParallelScheduler >>> from diffusers import StableDiffusionParadigmsPipeline - >>> scheduler = DDPMParallelScheduler.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="scheduler") + >>> scheduler = DDPMParallelScheduler.from_pretrained( + ... "stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="scheduler" + ... ) >>> pipe = StableDiffusionParadigmsPipeline.from_pretrained( - ... "runwayml/stable-diffusion-v1-5", scheduler=scheduler, torch_dtype=torch.float16 + ... "stable-diffusion-v1-5/stable-diffusion-v1-5", scheduler=scheduler, torch_dtype=torch.float16 ... ) >>> pipe = pipe.to("cuda") @@ -95,8 +97,8 @@ class StableDiffusionParadigmsPipeline( [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for + more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. """ diff --git a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_pix2pix_zero.py b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_pix2pix_zero.py index d000d87e6a..2a461ae20c 100644 --- a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_pix2pix_zero.py +++ b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_pix2pix_zero.py @@ -303,7 +303,8 @@ class StableDiffusionPix2PixZeroPipeline(DiffusionPipeline, StableDiffusionMixin [`DDIMScheduler`], [`LMSDiscreteScheduler`], [`EulerAncestralDiscreteScheduler`], or [`DDPMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details. + Please, refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for + details. feature_extractor ([`CLIPImageProcessor`]): Model that extracts features from generated images to be used as inputs for the `safety_checker`. requires_safety_checker (bool): diff --git a/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion.py b/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion.py index 61582853b0..9ff8e98577 100644 --- a/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion.py +++ b/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion.py @@ -38,8 +38,8 @@ class VersatileDiffusionPipeline(DiffusionPipeline): [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for + more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. """ diff --git a/src/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_img2img.py b/src/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_img2img.py index 1d7733982e..59f733a498 100644 --- a/src/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_img2img.py +++ b/src/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_img2img.py @@ -186,8 +186,8 @@ class LatentConsistencyModelImg2ImgPipeline( supports [`LCMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for + more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. requires_safety_checker (`bool`, *optional*, defaults to `True`): diff --git a/src/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_text2img.py b/src/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_text2img.py index 3e96b44663..e463884618 100644 --- a/src/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_text2img.py +++ b/src/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_text2img.py @@ -165,8 +165,8 @@ class LatentConsistencyModelPipeline( supports [`LCMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for + more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. requires_safety_checker (`bool`, *optional*, defaults to `True`): diff --git a/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py b/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py index 5b61aaf9b6..fbf4dc23d0 100644 --- a/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py +++ b/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py @@ -49,7 +49,7 @@ EXAMPLE_DOC_STRING = """ >>> from diffusers.utils import load_image >>> pipe = LEditsPPPipelineStableDiffusion.from_pretrained( - ... "runwayml/stable-diffusion-v1-5", variant="fp16", torch_dtype=torch.float16 + ... "stable-diffusion-v1-5/stable-diffusion-v1-5", variant="fp16", torch_dtype=torch.float16 ... ) >>> pipe.enable_vae_tiling() >>> pipe = pipe.to("cuda") @@ -381,8 +381,8 @@ class LEditsPPPipelineStableDiffusion( "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" + " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" diff --git a/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd.py b/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd.py index de66871922..1abef01430 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd.py @@ -80,7 +80,10 @@ EXAMPLE_DOC_STRING = """ >>> # load control net and stable diffusion v1-5 >>> controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16) >>> pipe = AutoPipelineForText2Image.from_pretrained( - ... "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16, enable_pag=True + ... "stable-diffusion-v1-5/stable-diffusion-v1-5", + ... controlnet=controlnet, + ... torch_dtype=torch.float16, + ... enable_pag=True, ... ) >>> # speed up diffusion process with faster scheduler and memory optimization @@ -202,8 +205,8 @@ class StableDiffusionControlNetPAGPipeline( [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for + more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. """ diff --git a/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_inpaint.py b/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_inpaint.py index 3daaac328c..2781af7890 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_inpaint.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_inpaint.py @@ -93,7 +93,10 @@ EXAMPLE_DOC_STRING = """ ... "lllyasviel/control_v11p_sd15_inpaint", torch_dtype=torch.float16 ... ) >>> pipe = AutoPipelineForInpainting.from_pretrained( - ... "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16, enable_pag=True + ... "stable-diffusion-v1-5/stable-diffusion-v1-5", + ... controlnet=controlnet, + ... torch_dtype=torch.float16, + ... enable_pag=True, ... ) >>> pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config) @@ -151,10 +154,11 @@ class StableDiffusionControlNetPAGInpaintPipeline( - [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters > [!TIP] > This pipeline can be used with checkpoints that have been specifically fine-tuned for inpainting > - ([runwayml/stable-diffusion-inpainting](https://huggingface.co/runwayml/stable-diffusion-inpainting)) as well as > - default text-to-image Stable Diffusion checkpoints > - ([runwayml/stable-diffusion-v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5)). Default text-to-image > - Stable Diffusion checkpoints might be preferable for ControlNets that have been fine-tuned on those, such as > + ([stable-diffusion-v1-5/stable-diffusion-inpainting](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-inpainting)) + as well as > default text-to-image Stable Diffusion checkpoints > + ([stable-diffusion-v1-5/stable-diffusion-v1-5](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5)). + Default text-to-image > Stable Diffusion checkpoints might be preferable for ControlNets that have been fine-tuned + on those, such as > [lllyasviel/control_v11p_sd15_inpaint](https://huggingface.co/lllyasviel/control_v11p_sd15_inpaint). Args: @@ -175,8 +179,8 @@ class StableDiffusionControlNetPAGInpaintPipeline( [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for + more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. """ @@ -1328,7 +1332,7 @@ class StableDiffusionControlNetPAGInpaintPipeline( # 7.1 Check that sizes of mask, masked image and latents match if num_channels_unet == 9: - # default case for runwayml/stable-diffusion-inpainting + # default case for stable-diffusion-v1-5/stable-diffusion-inpainting num_channels_mask = mask.shape[1] num_channels_masked_image = masked_image_latents.shape[1] if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels: diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd.py b/src/diffusers/pipelines/pag/pipeline_pag_sd.py index 349d006aad..ea64f8be2c 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd.py @@ -57,7 +57,7 @@ EXAMPLE_DOC_STRING = """ >>> from diffusers import AutoPipelineForText2Image >>> pipe = AutoPipelineForText2Image.from_pretrained( - ... "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, enable_pag=True + ... "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, enable_pag=True ... ) >>> pipe = pipe.to("cuda") @@ -190,8 +190,8 @@ class StableDiffusionPAGPipeline( [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for + more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. """ @@ -272,8 +272,8 @@ class StableDiffusionPAGPipeline( "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" + " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_img2img.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_img2img.py index e9a846b5e2..8351112ce4 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd_img2img.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_img2img.py @@ -61,7 +61,7 @@ EXAMPLE_DOC_STRING = """ >>> from diffusers.utils import load_image >>> pipe = AutoPipelineForImage2Image.from_pretrained( - ... "runwayml/stable-diffusion-v1-5", + ... "stable-diffusion-v1-5/stable-diffusion-v1-5", ... torch_dtype=torch.float16, ... enable_pag=True, ... ) @@ -185,8 +185,8 @@ class StableDiffusionPAGImg2ImgPipeline( [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for + more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. """ @@ -267,8 +267,8 @@ class StableDiffusionPAGImg2ImgPipeline( "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" + " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_inpaint.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_inpaint.py index ee9d20f368..6b1b294e10 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd_inpaint.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_inpaint.py @@ -58,7 +58,7 @@ EXAMPLE_DOC_STRING = """ >>> from diffusers import AutoPipelineForInpainting >>> pipe = AutoPipelineForInpainting.from_pretrained( - ... "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, enable_pag=True + ... "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, enable_pag=True ... ) >>> pipe = pipe.to("cuda") >>> img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png" @@ -217,8 +217,8 @@ class StableDiffusionPAGInpaintPipeline( [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for + more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. """ @@ -299,8 +299,8 @@ class StableDiffusionPAGInpaintPipeline( "The configuration file of the unet has set the default `sample_size` to smaller than" " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" + " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" " in the config might lead to incorrect results in future versions. If you have downloaded this" " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" @@ -1183,7 +1183,7 @@ class StableDiffusionPAGInpaintPipeline( # 8. Check that sizes of mask, masked image and latents match if num_channels_unet == 9: - # default case for runwayml/stable-diffusion-inpainting + # default case for stable-diffusion-v1-5/stable-diffusion-inpainting num_channels_mask = mask.shape[1] num_channels_masked_image = masked_image_latents.shape[1] if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels: diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_inpaint.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_inpaint.py index 2e12a4a97f..2a8f7a448d 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_inpaint.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_inpaint.py @@ -1501,7 +1501,7 @@ class StableDiffusionXLPAGInpaintPipeline( # 8. Check that sizes of mask, masked image and latents match if num_channels_unet == 9: - # default case for runwayml/stable-diffusion-inpainting + # default case for stable-diffusion-v1-5/stable-diffusion-inpainting num_channels_mask = mask.shape[1] num_channels_masked_image = masked_image_latents.shape[1] if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels: diff --git a/src/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py b/src/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py index 61435b80ca..c09992befb 100644 --- a/src/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py +++ b/src/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py @@ -179,8 +179,8 @@ class PaintByExamplePipeline(DeprecatedPipelineMixin, DiffusionPipeline, StableD [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for + more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. diff --git a/src/diffusers/pipelines/semantic_stable_diffusion/pipeline_semantic_stable_diffusion.py b/src/diffusers/pipelines/semantic_stable_diffusion/pipeline_semantic_stable_diffusion.py index a5f67bffe6..49b09e205c 100644 --- a/src/diffusers/pipelines/semantic_stable_diffusion/pipeline_semantic_stable_diffusion.py +++ b/src/diffusers/pipelines/semantic_stable_diffusion/pipeline_semantic_stable_diffusion.py @@ -48,8 +48,8 @@ class SemanticStableDiffusionPipeline(DeprecatedPipelineMixin, DiffusionPipeline [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`Q16SafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. + Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for + more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. """ @@ -332,7 +332,7 @@ class SemanticStableDiffusionPipeline(DeprecatedPipelineMixin, DiffusionPipeline >>> from diffusers import SemanticStableDiffusionPipeline >>> pipe = SemanticStableDiffusionPipeline.from_pretrained( - ... "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16 + ... "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16 ... ) >>> pipe = pipe.to("cuda") diff --git a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py index bed596e57c..5b2cca0378 100644 --- a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py +++ b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py @@ -1247,7 +1247,7 @@ class StableDiffusion3InpaintPipeline(DiffusionPipeline, SD3LoraLoaderMixin, Fro # match the inpainting pipeline and will be updated with input + mask inpainting model later if num_channels_transformer == 33: - # default case for runwayml/stable-diffusion-inpainting + # default case for stable-diffusion-v1-5/stable-diffusion-inpainting num_channels_mask = mask.shape[1] num_channels_masked_image = masked_image_latents.shape[1] if ( diff --git a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py index 18f8536a75..88cc7515b0 100644 --- a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py +++ b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py @@ -1501,7 +1501,7 @@ class StableDiffusionXLInpaintPipeline( # 8. Check that sizes of mask, masked image and latents match if num_channels_unet == 9: - # default case for runwayml/stable-diffusion-inpainting + # default case for stable-diffusion-v1-5/stable-diffusion-inpainting num_channels_mask = mask.shape[1] num_channels_masked_image = masked_image_latents.shape[1] if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels: diff --git a/src/diffusers/utils/hub_utils.py b/src/diffusers/utils/hub_utils.py index b6e99452aa..d0b05c7d95 100644 --- a/src/diffusers/utils/hub_utils.py +++ b/src/diffusers/utils/hub_utils.py @@ -113,7 +113,8 @@ def load_or_create_model_card( Args: repo_id_or_path (`str`): - The repo id (e.g., "runwayml/stable-diffusion-v1-5") or local path where to look for the model card. + The repo id (e.g., "stable-diffusion-v1-5/stable-diffusion-v1-5") or local path where to look for the model + card. token (`str`, *optional*): Authentication token. Will default to the stored token. See https://huggingface.co/settings/token for more details. diff --git a/tests/pipelines/pag/test_pag_sd_inpaint.py b/tests/pipelines/pag/test_pag_sd_inpaint.py index 709df68370..754158bbf1 100644 --- a/tests/pipelines/pag/test_pag_sd_inpaint.py +++ b/tests/pipelines/pag/test_pag_sd_inpaint.py @@ -255,7 +255,7 @@ class StableDiffusionPAGInpaintPipelineFastTests( @require_torch_accelerator class StableDiffusionPAGPipelineIntegrationTests(unittest.TestCase): pipeline_class = StableDiffusionPAGInpaintPipeline - repo_id = "runwayml/stable-diffusion-v1-5" + repo_id = "stable-diffusion-v1-5/stable-diffusion-v1-5" def setUp(self): super().setUp() From 23ebbb4bc81a17ebea17cb7cb94f301199e49a7f Mon Sep 17 00:00:00 2001 From: Lev Novitskiy <57654885+leffff@users.noreply.github.com> Date: Sat, 18 Oct 2025 07:34:30 +0300 Subject: [PATCH 074/123] Kandinsky 5 is finally in Diffusers! (#12478) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * add kandinsky5 transformer pipeline first version --------- Co-authored-by: Álvaro Somoza Co-authored-by: YiYi Xu Co-authored-by: Charles --- docs/source/en/api/loaders/lora.md | 3 + src/diffusers/__init__.py | 4 + src/diffusers/loaders/__init__.py | 2 + src/diffusers/loaders/lora_pipeline.py | 285 ++++++ src/diffusers/models/__init__.py | 2 + src/diffusers/models/transformers/__init__.py | 1 + .../transformers/transformer_kandinsky.py | 667 +++++++++++++ src/diffusers/pipelines/__init__.py | 2 + .../pipelines/kandinsky5/__init__.py | 48 + .../kandinsky5/pipeline_kandinsky.py | 893 ++++++++++++++++++ .../pipelines/kandinsky5/pipeline_output.py | 20 + src/diffusers/utils/dummy_pt_objects.py | 15 + .../dummy_torch_and_transformers_objects.py | 15 + 13 files changed, 1957 insertions(+) create mode 100644 src/diffusers/models/transformers/transformer_kandinsky.py create mode 100644 src/diffusers/pipelines/kandinsky5/__init__.py create mode 100644 src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py create mode 100644 src/diffusers/pipelines/kandinsky5/pipeline_output.py diff --git a/docs/source/en/api/loaders/lora.md b/docs/source/en/api/loaders/lora.md index b1d1ffb634..8e0326e0c3 100644 --- a/docs/source/en/api/loaders/lora.md +++ b/docs/source/en/api/loaders/lora.md @@ -107,6 +107,9 @@ LoRA is a fast and lightweight training method that inserts and trains a signifi [[autodoc]] loaders.lora_pipeline.QwenImageLoraLoaderMixin +## KandinskyLoraLoaderMixin +[[autodoc]] loaders.lora_pipeline.KandinskyLoraLoaderMixin + ## LoraBaseMixin [[autodoc]] loaders.lora_base.LoraBaseMixin \ No newline at end of file diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 95d559ff75..aa500b1494 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -220,6 +220,7 @@ else: "HunyuanVideoTransformer3DModel", "I2VGenXLUNet", "Kandinsky3UNet", + "Kandinsky5Transformer3DModel", "LatteTransformer3DModel", "LTXVideoTransformer3DModel", "Lumina2Transformer2DModel", @@ -474,6 +475,7 @@ else: "ImageTextPipelineOutput", "Kandinsky3Img2ImgPipeline", "Kandinsky3Pipeline", + "Kandinsky5T2VPipeline", "KandinskyCombinedPipeline", "KandinskyImg2ImgCombinedPipeline", "KandinskyImg2ImgPipeline", @@ -912,6 +914,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: HunyuanVideoTransformer3DModel, I2VGenXLUNet, Kandinsky3UNet, + Kandinsky5Transformer3DModel, LatteTransformer3DModel, LTXVideoTransformer3DModel, Lumina2Transformer2DModel, @@ -1136,6 +1139,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: ImageTextPipelineOutput, Kandinsky3Img2ImgPipeline, Kandinsky3Pipeline, + Kandinsky5T2VPipeline, KandinskyCombinedPipeline, KandinskyImg2ImgCombinedPipeline, KandinskyImg2ImgPipeline, diff --git a/src/diffusers/loaders/__init__.py b/src/diffusers/loaders/__init__.py index 7425486538..48507aae03 100644 --- a/src/diffusers/loaders/__init__.py +++ b/src/diffusers/loaders/__init__.py @@ -77,6 +77,7 @@ if is_torch_available(): "SanaLoraLoaderMixin", "Lumina2LoraLoaderMixin", "WanLoraLoaderMixin", + "KandinskyLoraLoaderMixin", "HiDreamImageLoraLoaderMixin", "SkyReelsV2LoraLoaderMixin", "QwenImageLoraLoaderMixin", @@ -115,6 +116,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: FluxLoraLoaderMixin, HiDreamImageLoraLoaderMixin, HunyuanVideoLoraLoaderMixin, + KandinskyLoraLoaderMixin, LoraLoaderMixin, LTXVideoLoraLoaderMixin, Lumina2LoraLoaderMixin, diff --git a/src/diffusers/loaders/lora_pipeline.py b/src/diffusers/loaders/lora_pipeline.py index e25a29e1c0..2bb6c0ea02 100644 --- a/src/diffusers/loaders/lora_pipeline.py +++ b/src/diffusers/loaders/lora_pipeline.py @@ -3639,6 +3639,291 @@ class Lumina2LoraLoaderMixin(LoraBaseMixin): super().unfuse_lora(components=components, **kwargs) +class KandinskyLoraLoaderMixin(LoraBaseMixin): + r""" + Load LoRA layers into [`Kandinsky5Transformer3DModel`], + """ + + _lora_loadable_modules = ["transformer"] + transformer_name = TRANSFORMER_NAME + + @classmethod + @validate_hf_hub_args + def lora_state_dict( + cls, + pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], + **kwargs, + ): + r""" + Return state dict for lora weights and the network alphas. + + Parameters: + pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): + Can be either: + - A string, the *model id* of a pretrained model hosted on the Hub. + - A path to a *directory* containing the model weights. + - A [torch state + dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). + + cache_dir (`Union[str, os.PathLike]`, *optional*): + Path to a directory where a downloaded pretrained model configuration is cached. + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force the (re-)download of the model weights. + proxies (`Dict[str, str]`, *optional*): + A dictionary of proxy servers to use by protocol or endpoint. + local_files_only (`bool`, *optional*, defaults to `False`): + Whether to only load local model weights and configuration files. + token (`str` or *bool*, *optional*): + The token to use as HTTP bearer authorization for remote files. + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. + subfolder (`str`, *optional*, defaults to `""`): + The subfolder location of a model file within a larger model repository. + weight_name (`str`, *optional*, defaults to None): + Name of the serialized state dict file. + use_safetensors (`bool`, *optional*): + Whether to use safetensors for loading. + return_lora_metadata (`bool`, *optional*, defaults to False): + When enabled, additionally return the LoRA adapter metadata. + """ + # Load the main state dict first which has the LoRA layers + cache_dir = kwargs.pop("cache_dir", None) + force_download = kwargs.pop("force_download", False) + proxies = kwargs.pop("proxies", None) + local_files_only = kwargs.pop("local_files_only", None) + token = kwargs.pop("token", None) + revision = kwargs.pop("revision", None) + subfolder = kwargs.pop("subfolder", None) + weight_name = kwargs.pop("weight_name", None) + use_safetensors = kwargs.pop("use_safetensors", None) + return_lora_metadata = kwargs.pop("return_lora_metadata", False) + + allow_pickle = False + if use_safetensors is None: + use_safetensors = True + allow_pickle = True + + user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"} + + state_dict, metadata = _fetch_state_dict( + pretrained_model_name_or_path_or_dict=pretrained_model_name_or_path_or_dict, + weight_name=weight_name, + use_safetensors=use_safetensors, + local_files_only=local_files_only, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + token=token, + revision=revision, + subfolder=subfolder, + user_agent=user_agent, + allow_pickle=allow_pickle, + ) + + is_dora_scale_present = any("dora_scale" in k for k in state_dict) + if is_dora_scale_present: + warn_msg = "It seems like you are using a DoRA checkpoint that is not compatible in Diffusers at the moment. So, we are going to filter out the keys associated to 'dora_scale` from the state dict. If you think this is a mistake please open an issue https://github.com/huggingface/diffusers/issues/new." + logger.warning(warn_msg) + state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k} + + out = (state_dict, metadata) if return_lora_metadata else state_dict + return out + + def load_lora_weights( + self, + pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], + adapter_name: Optional[str] = None, + hotswap: bool = False, + **kwargs, + ): + """ + Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` + + Parameters: + pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): + See [`~loaders.KandinskyLoraLoaderMixin.lora_state_dict`]. + adapter_name (`str`, *optional*): + Adapter name to be used for referencing the loaded adapter model. + hotswap (`bool`, *optional*): + Whether to substitute an existing (LoRA) adapter with the newly loaded adapter in-place. + low_cpu_mem_usage (`bool`, *optional*): + Speed up model loading by only loading the pretrained LoRA weights and not initializing the random + weights. + kwargs (`dict`, *optional*): + See [`~loaders.KandinskyLoraLoaderMixin.lora_state_dict`]. + """ + if not USE_PEFT_BACKEND: + raise ValueError("PEFT backend is required for this method.") + + low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT_LORA) + if low_cpu_mem_usage and not is_peft_version(">=", "0.13.1"): + raise ValueError( + "`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`." + ) + + # if a dict is passed, copy it instead of modifying it inplace + if isinstance(pretrained_model_name_or_path_or_dict, dict): + pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict.copy() + + # First, ensure that the checkpoint is a compatible one and can be successfully loaded. + kwargs["return_lora_metadata"] = True + state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) + + is_correct_format = all("lora" in key for key in state_dict.keys()) + if not is_correct_format: + raise ValueError("Invalid LoRA checkpoint.") + + # Load LoRA into transformer + self.load_lora_into_transformer( + state_dict, + transformer=getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer, + adapter_name=adapter_name, + metadata=metadata, + _pipeline=self, + low_cpu_mem_usage=low_cpu_mem_usage, + hotswap=hotswap, + ) + + @classmethod + def load_lora_into_transformer( + cls, + state_dict, + transformer, + adapter_name=None, + _pipeline=None, + low_cpu_mem_usage=False, + hotswap: bool = False, + metadata=None, + ): + """ + Load the LoRA layers specified in `state_dict` into `transformer`. + + Parameters: + state_dict (`dict`): + A standard state dict containing the lora layer parameters. + transformer (`Kandinsky5Transformer3DModel`): + The transformer model to load the LoRA layers into. + adapter_name (`str`, *optional*): + Adapter name to be used for referencing the loaded adapter model. + low_cpu_mem_usage (`bool`, *optional*): + Speed up model loading by only loading the pretrained LoRA weights. + hotswap (`bool`, *optional*): + See [`~loaders.KandinskyLoraLoaderMixin.load_lora_weights`]. + metadata (`dict`): + Optional LoRA adapter metadata. + """ + if low_cpu_mem_usage and not is_peft_version(">=", "0.13.1"): + raise ValueError( + "`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`." + ) + + # Load the layers corresponding to transformer. + logger.info(f"Loading {cls.transformer_name}.") + transformer.load_lora_adapter( + state_dict, + network_alphas=None, + adapter_name=adapter_name, + metadata=metadata, + _pipeline=_pipeline, + low_cpu_mem_usage=low_cpu_mem_usage, + hotswap=hotswap, + ) + + @classmethod + def save_lora_weights( + cls, + save_directory: Union[str, os.PathLike], + transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, + is_main_process: bool = True, + weight_name: str = None, + save_function: Callable = None, + safe_serialization: bool = True, + transformer_lora_adapter_metadata=None, + ): + r""" + Save the LoRA parameters corresponding to the transformer and text encoders. + + Arguments: + save_directory (`str` or `os.PathLike`): + Directory to save LoRA parameters to. + transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): + State dict of the LoRA layers corresponding to the `transformer`. + is_main_process (`bool`, *optional*, defaults to `True`): + Whether the process calling this is the main process. + save_function (`Callable`): + The function to use to save the state dictionary. + safe_serialization (`bool`, *optional*, defaults to `True`): + Whether to save the model using `safetensors` or the traditional PyTorch way. + transformer_lora_adapter_metadata: + LoRA adapter metadata associated with the transformer. + """ + lora_layers = {} + lora_metadata = {} + + if transformer_lora_layers: + lora_layers[cls.transformer_name] = transformer_lora_layers + lora_metadata[cls.transformer_name] = transformer_lora_adapter_metadata + + if not lora_layers: + raise ValueError("You must pass at least one of `transformer_lora_layers`") + + cls._save_lora_weights( + save_directory=save_directory, + lora_layers=lora_layers, + lora_metadata=lora_metadata, + is_main_process=is_main_process, + weight_name=weight_name, + save_function=save_function, + safe_serialization=safe_serialization, + ) + + def fuse_lora( + self, + components: List[str] = ["transformer"], + lora_scale: float = 1.0, + safe_fusing: bool = False, + adapter_names: Optional[List[str]] = None, + **kwargs, + ): + r""" + Fuses the LoRA parameters into the original parameters of the corresponding blocks. + + Args: + components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into. + lora_scale (`float`, defaults to 1.0): + Controls how much to influence the outputs with the LoRA parameters. + safe_fusing (`bool`, defaults to `False`): + Whether to check fused weights for NaN values before fusing. + adapter_names (`List[str]`, *optional*): + Adapter names to be used for fusing. + + Example: + ```py + from diffusers import Kandinsky5T2VPipeline + + pipeline = Kandinsky5T2VPipeline.from_pretrained("ai-forever/Kandinsky-5.0-T2V") + pipeline.load_lora_weights("path/to/lora.safetensors") + pipeline.fuse_lora(lora_scale=0.7) + ``` + """ + super().fuse_lora( + components=components, + lora_scale=lora_scale, + safe_fusing=safe_fusing, + adapter_names=adapter_names, + **kwargs, + ) + + def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): + r""" + Reverses the effect of [`pipe.fuse_lora()`]. + + Args: + components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. + """ + super().unfuse_lora(components=components, **kwargs) + + class WanLoraLoaderMixin(LoraBaseMixin): r""" Load LoRA layers into [`WanTransformer3DModel`]. Specific to [`WanPipeline`] and `[WanImageToVideoPipeline`]. diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index 457f70448a..8d029bf5d3 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -91,6 +91,7 @@ if is_torch_available(): _import_structure["transformers.transformer_hidream_image"] = ["HiDreamImageTransformer2DModel"] _import_structure["transformers.transformer_hunyuan_video"] = ["HunyuanVideoTransformer3DModel"] _import_structure["transformers.transformer_hunyuan_video_framepack"] = ["HunyuanVideoFramepackTransformer3DModel"] + _import_structure["transformers.transformer_kandinsky"] = ["Kandinsky5Transformer3DModel"] _import_structure["transformers.transformer_ltx"] = ["LTXVideoTransformer3DModel"] _import_structure["transformers.transformer_lumina2"] = ["Lumina2Transformer2DModel"] _import_structure["transformers.transformer_mochi"] = ["MochiTransformer3DModel"] @@ -182,6 +183,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: HunyuanDiT2DModel, HunyuanVideoFramepackTransformer3DModel, HunyuanVideoTransformer3DModel, + Kandinsky5Transformer3DModel, LatteTransformer3DModel, LTXVideoTransformer3DModel, Lumina2Transformer2DModel, diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index b60f0636e6..6b80ea6c82 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -27,6 +27,7 @@ if is_torch_available(): from .transformer_hidream_image import HiDreamImageTransformer2DModel from .transformer_hunyuan_video import HunyuanVideoTransformer3DModel from .transformer_hunyuan_video_framepack import HunyuanVideoFramepackTransformer3DModel + from .transformer_kandinsky import Kandinsky5Transformer3DModel from .transformer_ltx import LTXVideoTransformer3DModel from .transformer_lumina2 import Lumina2Transformer2DModel from .transformer_mochi import MochiTransformer3DModel diff --git a/src/diffusers/models/transformers/transformer_kandinsky.py b/src/diffusers/models/transformers/transformer_kandinsky.py new file mode 100644 index 0000000000..d4ba92acaf --- /dev/null +++ b/src/diffusers/models/transformers/transformer_kandinsky.py @@ -0,0 +1,667 @@ +# Copyright 2025 The Kandinsky Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import math +from typing import Any, Dict, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin, PeftAdapterMixin +from ...utils import ( + logging, +) +from ..attention import AttentionMixin, AttentionModuleMixin +from ..attention_dispatch import _CAN_USE_FLEX_ATTN, dispatch_attention_fn +from ..cache_utils import CacheMixin +from ..modeling_outputs import Transformer2DModelOutput +from ..modeling_utils import ModelMixin + + +logger = logging.get_logger(__name__) + + +def get_freqs(dim, max_period=10000.0): + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=dim, dtype=torch.float32) / dim) + return freqs + + +def fractal_flatten(x, rope, shape, block_mask=False): + if block_mask: + pixel_size = 8 + x = local_patching(x, shape, (1, pixel_size, pixel_size), dim=1) + rope = local_patching(rope, shape, (1, pixel_size, pixel_size), dim=1) + x = x.flatten(1, 2) + rope = rope.flatten(1, 2) + else: + x = x.flatten(1, 3) + rope = rope.flatten(1, 3) + return x, rope + + +def fractal_unflatten(x, shape, block_mask=False): + if block_mask: + pixel_size = 8 + x = x.reshape(x.shape[0], -1, pixel_size**2, *x.shape[2:]) + x = local_merge(x, shape, (1, pixel_size, pixel_size), dim=1) + else: + x = x.reshape(*shape, *x.shape[2:]) + return x + + +def local_patching(x, shape, group_size, dim=0): + batch_size, duration, height, width = shape + g1, g2, g3 = group_size + x = x.reshape( + *x.shape[:dim], + duration // g1, + g1, + height // g2, + g2, + width // g3, + g3, + *x.shape[dim + 3 :], + ) + x = x.permute( + *range(len(x.shape[:dim])), + dim, + dim + 2, + dim + 4, + dim + 1, + dim + 3, + dim + 5, + *range(dim + 6, len(x.shape)), + ) + x = x.flatten(dim, dim + 2).flatten(dim + 1, dim + 3) + return x + + +def local_merge(x, shape, group_size, dim=0): + batch_size, duration, height, width = shape + g1, g2, g3 = group_size + x = x.reshape( + *x.shape[:dim], + duration // g1, + height // g2, + width // g3, + g1, + g2, + g3, + *x.shape[dim + 2 :], + ) + x = x.permute( + *range(len(x.shape[:dim])), + dim, + dim + 3, + dim + 1, + dim + 4, + dim + 2, + dim + 5, + *range(dim + 6, len(x.shape)), + ) + x = x.flatten(dim, dim + 1).flatten(dim + 1, dim + 2).flatten(dim + 2, dim + 3) + return x + + +def nablaT_v2( + q: Tensor, + k: Tensor, + sta: Tensor, + thr: float = 0.9, +): + if _CAN_USE_FLEX_ATTN: + from torch.nn.attention.flex_attention import BlockMask + else: + raise ValueError("Nabla attention is not supported with this version of PyTorch") + + q = q.transpose(1, 2).contiguous() + k = k.transpose(1, 2).contiguous() + + # Map estimation + B, h, S, D = q.shape + s1 = S // 64 + qa = q.reshape(B, h, s1, 64, D).mean(-2) + ka = k.reshape(B, h, s1, 64, D).mean(-2).transpose(-2, -1) + map = qa @ ka + + map = torch.softmax(map / math.sqrt(D), dim=-1) + # Map binarization + vals, inds = map.sort(-1) + cvals = vals.cumsum_(-1) + mask = (cvals >= 1 - thr).int() + mask = mask.gather(-1, inds.argsort(-1)) + + mask = torch.logical_or(mask, sta) + + # BlockMask creation + kv_nb = mask.sum(-1).to(torch.int32) + kv_inds = mask.argsort(dim=-1, descending=True).to(torch.int32) + return BlockMask.from_kv_blocks(torch.zeros_like(kv_nb), kv_inds, kv_nb, kv_inds, BLOCK_SIZE=64, mask_mod=None) + + +class Kandinsky5TimeEmbeddings(nn.Module): + def __init__(self, model_dim, time_dim, max_period=10000.0): + super().__init__() + assert model_dim % 2 == 0 + self.model_dim = model_dim + self.max_period = max_period + self.freqs = get_freqs(self.model_dim // 2, self.max_period) + self.in_layer = nn.Linear(model_dim, time_dim, bias=True) + self.activation = nn.SiLU() + self.out_layer = nn.Linear(time_dim, time_dim, bias=True) + + @torch.autocast(device_type="cuda", dtype=torch.float32) + def forward(self, time): + args = torch.outer(time, self.freqs.to(device=time.device)) + time_embed = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + time_embed = self.out_layer(self.activation(self.in_layer(time_embed))) + return time_embed + + +class Kandinsky5TextEmbeddings(nn.Module): + def __init__(self, text_dim, model_dim): + super().__init__() + self.in_layer = nn.Linear(text_dim, model_dim, bias=True) + self.norm = nn.LayerNorm(model_dim, elementwise_affine=True) + + def forward(self, text_embed): + text_embed = self.in_layer(text_embed) + return self.norm(text_embed).type_as(text_embed) + + +class Kandinsky5VisualEmbeddings(nn.Module): + def __init__(self, visual_dim, model_dim, patch_size): + super().__init__() + self.patch_size = patch_size + self.in_layer = nn.Linear(math.prod(patch_size) * visual_dim, model_dim) + + def forward(self, x): + batch_size, duration, height, width, dim = x.shape + x = ( + x.view( + batch_size, + duration // self.patch_size[0], + self.patch_size[0], + height // self.patch_size[1], + self.patch_size[1], + width // self.patch_size[2], + self.patch_size[2], + dim, + ) + .permute(0, 1, 3, 5, 2, 4, 6, 7) + .flatten(4, 7) + ) + return self.in_layer(x) + + +class Kandinsky5RoPE1D(nn.Module): + def __init__(self, dim, max_pos=1024, max_period=10000.0): + super().__init__() + self.max_period = max_period + self.dim = dim + self.max_pos = max_pos + freq = get_freqs(dim // 2, max_period) + pos = torch.arange(max_pos, dtype=freq.dtype) + self.register_buffer("args", torch.outer(pos, freq), persistent=False) + + def forward(self, pos): + args = self.args[pos] + cosine = torch.cos(args) + sine = torch.sin(args) + rope = torch.stack([cosine, -sine, sine, cosine], dim=-1) + rope = rope.view(*rope.shape[:-1], 2, 2) + return rope.unsqueeze(-4) + + +class Kandinsky5RoPE3D(nn.Module): + def __init__(self, axes_dims, max_pos=(128, 128, 128), max_period=10000.0): + super().__init__() + self.axes_dims = axes_dims + self.max_pos = max_pos + self.max_period = max_period + + for i, (axes_dim, ax_max_pos) in enumerate(zip(axes_dims, max_pos)): + freq = get_freqs(axes_dim // 2, max_period) + pos = torch.arange(ax_max_pos, dtype=freq.dtype) + self.register_buffer(f"args_{i}", torch.outer(pos, freq), persistent=False) + + def forward(self, shape, pos, scale_factor=(1.0, 1.0, 1.0)): + batch_size, duration, height, width = shape + args_t = self.args_0[pos[0]] / scale_factor[0] + args_h = self.args_1[pos[1]] / scale_factor[1] + args_w = self.args_2[pos[2]] / scale_factor[2] + + args = torch.cat( + [ + args_t.view(1, duration, 1, 1, -1).repeat(batch_size, 1, height, width, 1), + args_h.view(1, 1, height, 1, -1).repeat(batch_size, duration, 1, width, 1), + args_w.view(1, 1, 1, width, -1).repeat(batch_size, duration, height, 1, 1), + ], + dim=-1, + ) + cosine = torch.cos(args) + sine = torch.sin(args) + rope = torch.stack([cosine, -sine, sine, cosine], dim=-1) + rope = rope.view(*rope.shape[:-1], 2, 2) + return rope.unsqueeze(-4) + + +class Kandinsky5Modulation(nn.Module): + def __init__(self, time_dim, model_dim, num_params): + super().__init__() + self.activation = nn.SiLU() + self.out_layer = nn.Linear(time_dim, num_params * model_dim) + self.out_layer.weight.data.zero_() + self.out_layer.bias.data.zero_() + + @torch.autocast(device_type="cuda", dtype=torch.float32) + def forward(self, x): + return self.out_layer(self.activation(x)) + + +class Kandinsky5AttnProcessor: + _attention_backend = None + _parallel_config = None + + def __init__(self): + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError(f"{self.__class__.__name__} requires PyTorch 2.0. Please upgrade your pytorch version.") + + def __call__(self, attn, hidden_states, encoder_hidden_states=None, rotary_emb=None, sparse_params=None): + # query, key, value = self.get_qkv(x) + query = attn.to_query(hidden_states) + + if encoder_hidden_states is not None: + key = attn.to_key(encoder_hidden_states) + value = attn.to_value(encoder_hidden_states) + + shape, cond_shape = query.shape[:-1], key.shape[:-1] + query = query.reshape(*shape, attn.num_heads, -1) + key = key.reshape(*cond_shape, attn.num_heads, -1) + value = value.reshape(*cond_shape, attn.num_heads, -1) + + else: + key = attn.to_key(hidden_states) + value = attn.to_value(hidden_states) + + shape = query.shape[:-1] + query = query.reshape(*shape, attn.num_heads, -1) + key = key.reshape(*shape, attn.num_heads, -1) + value = value.reshape(*shape, attn.num_heads, -1) + + # query, key = self.norm_qk(query, key) + query = attn.query_norm(query.float()).type_as(query) + key = attn.key_norm(key.float()).type_as(key) + + def apply_rotary(x, rope): + x_ = x.reshape(*x.shape[:-1], -1, 1, 2).to(torch.float32) + x_out = (rope * x_).sum(dim=-1) + return x_out.reshape(*x.shape).to(torch.bfloat16) + + if rotary_emb is not None: + query = apply_rotary(query, rotary_emb).type_as(query) + key = apply_rotary(key, rotary_emb).type_as(key) + + if sparse_params is not None: + attn_mask = nablaT_v2( + query, + key, + sparse_params["sta_mask"], + thr=sparse_params["P"], + ) + else: + attn_mask = None + + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=attn_mask, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + hidden_states = hidden_states.flatten(-2, -1) + + attn_out = attn.out_layer(hidden_states) + return attn_out + + +class Kandinsky5Attention(nn.Module, AttentionModuleMixin): + _default_processor_cls = Kandinsky5AttnProcessor + _available_processors = [ + Kandinsky5AttnProcessor, + ] + + def __init__(self, num_channels, head_dim, processor=None): + super().__init__() + assert num_channels % head_dim == 0 + self.num_heads = num_channels // head_dim + + self.to_query = nn.Linear(num_channels, num_channels, bias=True) + self.to_key = nn.Linear(num_channels, num_channels, bias=True) + self.to_value = nn.Linear(num_channels, num_channels, bias=True) + self.query_norm = nn.RMSNorm(head_dim) + self.key_norm = nn.RMSNorm(head_dim) + + self.out_layer = nn.Linear(num_channels, num_channels, bias=True) + if processor is None: + processor = self._default_processor_cls() + self.set_processor(processor) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + sparse_params: Optional[torch.Tensor] = None, + rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + **kwargs, + ) -> torch.Tensor: + attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys()) + quiet_attn_parameters = {} + unused_kwargs = [k for k, _ in kwargs.items() if k not in attn_parameters and k not in quiet_attn_parameters] + if len(unused_kwargs) > 0: + logger.warning( + f"attention_processor_kwargs {unused_kwargs} are not expected by {self.processor.__class__.__name__} and will be ignored." + ) + kwargs = {k: w for k, w in kwargs.items() if k in attn_parameters} + + return self.processor( + self, + hidden_states, + encoder_hidden_states=encoder_hidden_states, + sparse_params=sparse_params, + rotary_emb=rotary_emb, + **kwargs, + ) + + +class Kandinsky5FeedForward(nn.Module): + def __init__(self, dim, ff_dim): + super().__init__() + self.in_layer = nn.Linear(dim, ff_dim, bias=False) + self.activation = nn.GELU() + self.out_layer = nn.Linear(ff_dim, dim, bias=False) + + def forward(self, x): + return self.out_layer(self.activation(self.in_layer(x))) + + +class Kandinsky5OutLayer(nn.Module): + def __init__(self, model_dim, time_dim, visual_dim, patch_size): + super().__init__() + self.patch_size = patch_size + self.modulation = Kandinsky5Modulation(time_dim, model_dim, 2) + self.norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.out_layer = nn.Linear(model_dim, math.prod(patch_size) * visual_dim, bias=True) + + def forward(self, visual_embed, text_embed, time_embed): + shift, scale = torch.chunk(self.modulation(time_embed).unsqueeze(dim=1), 2, dim=-1) + + visual_embed = ( + self.norm(visual_embed.float()) * (scale.float()[:, None, None] + 1.0) + shift.float()[:, None, None] + ).type_as(visual_embed) + + x = self.out_layer(visual_embed) + + batch_size, duration, height, width, _ = x.shape + x = ( + x.view( + batch_size, + duration, + height, + width, + -1, + self.patch_size[0], + self.patch_size[1], + self.patch_size[2], + ) + .permute(0, 1, 5, 2, 6, 3, 7, 4) + .flatten(1, 2) + .flatten(2, 3) + .flatten(3, 4) + ) + return x + + +class Kandinsky5TransformerEncoderBlock(nn.Module): + def __init__(self, model_dim, time_dim, ff_dim, head_dim): + super().__init__() + self.text_modulation = Kandinsky5Modulation(time_dim, model_dim, 6) + + self.self_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.self_attention = Kandinsky5Attention(model_dim, head_dim, processor=Kandinsky5AttnProcessor()) + + self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.feed_forward = Kandinsky5FeedForward(model_dim, ff_dim) + + def forward(self, x, time_embed, rope): + self_attn_params, ff_params = torch.chunk(self.text_modulation(time_embed).unsqueeze(dim=1), 2, dim=-1) + shift, scale, gate = torch.chunk(self_attn_params, 3, dim=-1) + out = (self.self_attention_norm(x.float()) * (scale.float() + 1.0) + shift.float()).type_as(x) + out = self.self_attention(out, rotary_emb=rope) + x = (x.float() + gate.float() * out.float()).type_as(x) + + shift, scale, gate = torch.chunk(ff_params, 3, dim=-1) + out = (self.feed_forward_norm(x.float()) * (scale.float() + 1.0) + shift.float()).type_as(x) + out = self.feed_forward(out) + x = (x.float() + gate.float() * out.float()).type_as(x) + + return x + + +class Kandinsky5TransformerDecoderBlock(nn.Module): + def __init__(self, model_dim, time_dim, ff_dim, head_dim): + super().__init__() + self.visual_modulation = Kandinsky5Modulation(time_dim, model_dim, 9) + + self.self_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.self_attention = Kandinsky5Attention(model_dim, head_dim, processor=Kandinsky5AttnProcessor()) + + self.cross_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.cross_attention = Kandinsky5Attention(model_dim, head_dim, processor=Kandinsky5AttnProcessor()) + + self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.feed_forward = Kandinsky5FeedForward(model_dim, ff_dim) + + def forward(self, visual_embed, text_embed, time_embed, rope, sparse_params): + self_attn_params, cross_attn_params, ff_params = torch.chunk( + self.visual_modulation(time_embed).unsqueeze(dim=1), 3, dim=-1 + ) + + shift, scale, gate = torch.chunk(self_attn_params, 3, dim=-1) + visual_out = (self.self_attention_norm(visual_embed.float()) * (scale.float() + 1.0) + shift.float()).type_as( + visual_embed + ) + visual_out = self.self_attention(visual_out, rotary_emb=rope, sparse_params=sparse_params) + visual_embed = (visual_embed.float() + gate.float() * visual_out.float()).type_as(visual_embed) + + shift, scale, gate = torch.chunk(cross_attn_params, 3, dim=-1) + visual_out = (self.cross_attention_norm(visual_embed.float()) * (scale.float() + 1.0) + shift.float()).type_as( + visual_embed + ) + visual_out = self.cross_attention(visual_out, encoder_hidden_states=text_embed) + visual_embed = (visual_embed.float() + gate.float() * visual_out.float()).type_as(visual_embed) + + shift, scale, gate = torch.chunk(ff_params, 3, dim=-1) + visual_out = (self.feed_forward_norm(visual_embed.float()) * (scale.float() + 1.0) + shift.float()).type_as( + visual_embed + ) + visual_out = self.feed_forward(visual_out) + visual_embed = (visual_embed.float() + gate.float() * visual_out.float()).type_as(visual_embed) + + return visual_embed + + +class Kandinsky5Transformer3DModel( + ModelMixin, + ConfigMixin, + PeftAdapterMixin, + FromOriginalModelMixin, + CacheMixin, + AttentionMixin, +): + """ + A 3D Diffusion Transformer model for video-like data. + """ + + _repeated_blocks = [ + "Kandinsky5TransformerEncoderBlock", + "Kandinsky5TransformerDecoderBlock", + ] + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + in_visual_dim=4, + in_text_dim=3584, + in_text_dim2=768, + time_dim=512, + out_visual_dim=4, + patch_size=(1, 2, 2), + model_dim=2048, + ff_dim=5120, + num_text_blocks=2, + num_visual_blocks=32, + axes_dims=(16, 24, 24), + visual_cond=False, + attention_type: str = "regular", + attention_causal: bool = None, + attention_local: bool = None, + attention_glob: bool = None, + attention_window: int = None, + attention_P: float = None, + attention_wT: int = None, + attention_wW: int = None, + attention_wH: int = None, + attention_add_sta: bool = None, + attention_method: str = None, + ): + super().__init__() + + head_dim = sum(axes_dims) + self.in_visual_dim = in_visual_dim + self.model_dim = model_dim + self.patch_size = patch_size + self.visual_cond = visual_cond + self.attention_type = attention_type + + visual_embed_dim = 2 * in_visual_dim + 1 if visual_cond else in_visual_dim + + # Initialize embeddings + self.time_embeddings = Kandinsky5TimeEmbeddings(model_dim, time_dim) + self.text_embeddings = Kandinsky5TextEmbeddings(in_text_dim, model_dim) + self.pooled_text_embeddings = Kandinsky5TextEmbeddings(in_text_dim2, time_dim) + self.visual_embeddings = Kandinsky5VisualEmbeddings(visual_embed_dim, model_dim, patch_size) + + # Initialize positional embeddings + self.text_rope_embeddings = Kandinsky5RoPE1D(head_dim) + self.visual_rope_embeddings = Kandinsky5RoPE3D(axes_dims) + + # Initialize transformer blocks + self.text_transformer_blocks = nn.ModuleList( + [Kandinsky5TransformerEncoderBlock(model_dim, time_dim, ff_dim, head_dim) for _ in range(num_text_blocks)] + ) + + self.visual_transformer_blocks = nn.ModuleList( + [ + Kandinsky5TransformerDecoderBlock(model_dim, time_dim, ff_dim, head_dim) + for _ in range(num_visual_blocks) + ] + ) + + # Initialize output layer + self.out_layer = Kandinsky5OutLayer(model_dim, time_dim, out_visual_dim, patch_size) + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, # x + encoder_hidden_states: torch.Tensor, # text_embed + timestep: torch.Tensor, # time + pooled_projections: torch.Tensor, # pooled_text_embed + visual_rope_pos: Tuple[int, int, int], + text_rope_pos: torch.LongTensor, + scale_factor: Tuple[float, float, float] = (1.0, 1.0, 1.0), + sparse_params: Optional[Dict[str, Any]] = None, + return_dict: bool = True, + ) -> Union[Transformer2DModelOutput, torch.FloatTensor]: + """ + Forward pass of the Kandinsky5 3D Transformer. + + Args: + hidden_states (`torch.FloatTensor`): Input visual states + encoder_hidden_states (`torch.FloatTensor`): Text embeddings + timestep (`torch.Tensor` or `float` or `int`): Current timestep + pooled_projections (`torch.FloatTensor`): Pooled text embeddings + visual_rope_pos (`Tuple[int, int, int]`): Position for visual RoPE + text_rope_pos (`torch.LongTensor`): Position for text RoPE + scale_factor (`Tuple[float, float, float]`, optional): Scale factor for RoPE + sparse_params (`Dict[str, Any]`, optional): Parameters for sparse attention + return_dict (`bool`, optional): Whether to return a dictionary + + Returns: + [`~models.transformer_2d.Transformer2DModelOutput`] or `torch.FloatTensor`: The output of the transformer + """ + x = hidden_states + text_embed = encoder_hidden_states + time = timestep + pooled_text_embed = pooled_projections + + text_embed = self.text_embeddings(text_embed) + time_embed = self.time_embeddings(time) + time_embed = time_embed + self.pooled_text_embeddings(pooled_text_embed) + visual_embed = self.visual_embeddings(x) + text_rope = self.text_rope_embeddings(text_rope_pos) + text_rope = text_rope.unsqueeze(dim=0) + + for text_transformer_block in self.text_transformer_blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + text_embed = self._gradient_checkpointing_func( + text_transformer_block, text_embed, time_embed, text_rope + ) + else: + text_embed = text_transformer_block(text_embed, time_embed, text_rope) + + visual_shape = visual_embed.shape[:-1] + visual_rope = self.visual_rope_embeddings(visual_shape, visual_rope_pos, scale_factor) + to_fractal = sparse_params["to_fractal"] if sparse_params is not None else False + visual_embed, visual_rope = fractal_flatten(visual_embed, visual_rope, visual_shape, block_mask=to_fractal) + + for visual_transformer_block in self.visual_transformer_blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + visual_embed = self._gradient_checkpointing_func( + visual_transformer_block, + visual_embed, + text_embed, + time_embed, + visual_rope, + sparse_params, + ) + else: + visual_embed = visual_transformer_block( + visual_embed, text_embed, time_embed, visual_rope, sparse_params + ) + + visual_embed = fractal_unflatten(visual_embed, visual_shape, block_mask=to_fractal) + x = self.out_layer(visual_embed, text_embed, time_embed) + + if not return_dict: + return x + + return Transformer2DModelOutput(sample=x) diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index 190c7871d2..c438caed57 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -382,6 +382,7 @@ else: "WuerstchenPriorPipeline", ] _import_structure["wan"] = ["WanPipeline", "WanImageToVideoPipeline", "WanVideoToVideoPipeline", "WanVACEPipeline"] + _import_structure["kandinsky5"] = ["Kandinsky5T2VPipeline"] _import_structure["skyreels_v2"] = [ "SkyReelsV2DiffusionForcingPipeline", "SkyReelsV2DiffusionForcingImageToVideoPipeline", @@ -671,6 +672,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: Kandinsky3Img2ImgPipeline, Kandinsky3Pipeline, ) + from .kandinsky5 import Kandinsky5T2VPipeline from .latent_consistency_models import ( LatentConsistencyModelImg2ImgPipeline, LatentConsistencyModelPipeline, diff --git a/src/diffusers/pipelines/kandinsky5/__init__.py b/src/diffusers/pipelines/kandinsky5/__init__.py new file mode 100644 index 0000000000..a7975bdce9 --- /dev/null +++ b/src/diffusers/pipelines/kandinsky5/__init__.py @@ -0,0 +1,48 @@ +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects # noqa F403 + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["pipeline_kandinsky"] = ["Kandinsky5T2VPipeline"] + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() + + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * + else: + from .pipeline_kandinsky import Kandinsky5T2VPipeline + +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) + + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py new file mode 100644 index 0000000000..a1122a8256 --- /dev/null +++ b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py @@ -0,0 +1,893 @@ +# Copyright 2025 The Kandinsky Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import html +from typing import Callable, Dict, List, Optional, Union + +import regex as re +import torch +from torch.nn import functional as F +from transformers import CLIPTextModel, CLIPTokenizer, Qwen2_5_VLForConditionalGeneration, Qwen2VLProcessor + +from ...callbacks import MultiPipelineCallbacks, PipelineCallback +from ...loaders import KandinskyLoraLoaderMixin +from ...models import AutoencoderKLHunyuanVideo +from ...models.transformers import Kandinsky5Transformer3DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import is_ftfy_available, is_torch_xla_available, logging, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ...video_processor import VideoProcessor +from ..pipeline_utils import DiffusionPipeline +from .pipeline_output import KandinskyPipelineOutput + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +if is_ftfy_available(): + import ftfy + + +logger = logging.get_logger(__name__) + +EXAMPLE_DOC_STRING = """ + Examples: + + ```python + >>> import torch + >>> from diffusers import Kandinsky5T2VPipeline + >>> from diffusers.utils import export_to_video + + >>> # Available models: + >>> # ai-forever/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers + >>> # ai-forever/Kandinsky-5.0-T2V-Lite-nocfg-5s-Diffusers + >>> # ai-forever/Kandinsky-5.0-T2V-Lite-distilled16steps-5s-Diffusers + >>> # ai-forever/Kandinsky-5.0-T2V-Lite-pretrain-5s-Diffusers + + >>> model_id = "ai-forever/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers" + >>> pipe = Kandinsky5T2VPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16) + >>> pipe = pipe.to("cuda") + + >>> prompt = "A cat and a dog baking a cake together in a kitchen." + >>> negative_prompt = "Static, 2D cartoon, cartoon, 2d animation, paintings, images, worst quality, low quality, ugly, deformed, walking backwards" + + >>> output = pipe( + ... prompt=prompt, + ... negative_prompt=negative_prompt, + ... height=512, + ... width=768, + ... num_frames=121, + ... num_inference_steps=50, + ... guidance_scale=5.0, + ... ).frames[0] + + >>> export_to_video(output, "output.mp4", fps=24, quality=9) + ``` +""" + + +def basic_clean(text): + """Clean text using ftfy if available and unescape HTML entities.""" + if is_ftfy_available(): + text = ftfy.fix_text(text) + text = html.unescape(html.unescape(text)) + return text.strip() + + +def whitespace_clean(text): + """Normalize whitespace in text by replacing multiple spaces with single space.""" + text = re.sub(r"\s+", " ", text) + text = text.strip() + return text + + +def prompt_clean(text): + """Apply both basic cleaning and whitespace normalization to prompts.""" + text = whitespace_clean(basic_clean(text)) + return text + + +class Kandinsky5T2VPipeline(DiffusionPipeline, KandinskyLoraLoaderMixin): + r""" + Pipeline for text-to-video generation using Kandinsky 5.0. + + This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods + implemented for all pipelines (downloading, saving, running on a particular device, etc.). + + Args: + transformer ([`Kandinsky5Transformer3DModel`]): + Conditional Transformer to denoise the encoded video latents. + vae ([`AutoencoderKLHunyuanVideo`]): + Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations. + text_encoder ([`Qwen2_5_VLForConditionalGeneration`]): + Frozen text-encoder (Qwen2.5-VL). + tokenizer ([`AutoProcessor`]): + Tokenizer for Qwen2.5-VL. + text_encoder_2 ([`CLIPTextModel`]): + Frozen CLIP text encoder. + tokenizer_2 ([`CLIPTokenizer`]): + Tokenizer for CLIP. + scheduler ([`FlowMatchEulerDiscreteScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded video latents. + """ + + model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae" + _callback_tensor_inputs = [ + "latents", + "prompt_embeds_qwen", + "prompt_embeds_clip", + "negative_prompt_embeds_qwen", + "negative_prompt_embeds_clip", + ] + + def __init__( + self, + transformer: Kandinsky5Transformer3DModel, + vae: AutoencoderKLHunyuanVideo, + text_encoder: Qwen2_5_VLForConditionalGeneration, + tokenizer: Qwen2VLProcessor, + text_encoder_2: CLIPTextModel, + tokenizer_2: CLIPTokenizer, + scheduler: FlowMatchEulerDiscreteScheduler, + ): + super().__init__() + + self.register_modules( + transformer=transformer, + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + text_encoder_2=text_encoder_2, + tokenizer_2=tokenizer_2, + scheduler=scheduler, + ) + + self.prompt_template = "\n".join( + [ + "<|im_start|>system\nYou are a promt engineer. Describe the video in detail.", + "Describe how the camera moves or shakes, describe the zoom and view angle, whether it follows the objects.", + "Describe the location of the video, main characters or objects and their action.", + "Describe the dynamism of the video and presented actions.", + "Name the visual style of the video: whether it is a professional footage, user generated content, some kind of animation, video game or scren content.", + "Describe the visual effects, postprocessing and transitions if they are presented in the video.", + "Pay attention to the order of key actions shown in the scene.<|im_end|>", + "<|im_start|>user\n{}<|im_end|>", + ] + ) + self.prompt_template_encode_start_idx = 129 + + self.vae_scale_factor_temporal = vae.config.temporal_compression_ratio + self.vae_scale_factor_spatial = vae.config.spatial_compression_ratio + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) + + @staticmethod + def fast_sta_nabla(T: int, H: int, W: int, wT: int = 3, wH: int = 3, wW: int = 3, device="cuda") -> torch.Tensor: + """ + Create a sparse temporal attention (STA) mask for efficient video generation. + + This method generates a mask that limits attention to nearby frames and spatial positions, reducing + computational complexity for video generation. + + Args: + T (int): Number of temporal frames + H (int): Height in latent space + W (int): Width in latent space + wT (int): Temporal attention window size + wH (int): Height attention window size + wW (int): Width attention window size + device (str): Device to create tensor on + + Returns: + torch.Tensor: Sparse attention mask of shape (T*H*W, T*H*W) + """ + l = torch.Tensor([T, H, W]).amax() + r = torch.arange(0, l, 1, dtype=torch.int16, device=device) + mat = (r.unsqueeze(1) - r.unsqueeze(0)).abs() + sta_t, sta_h, sta_w = ( + mat[:T, :T].flatten(), + mat[:H, :H].flatten(), + mat[:W, :W].flatten(), + ) + sta_t = sta_t <= wT // 2 + sta_h = sta_h <= wH // 2 + sta_w = sta_w <= wW // 2 + sta_hw = (sta_h.unsqueeze(1) * sta_w.unsqueeze(0)).reshape(H, H, W, W).transpose(1, 2).flatten() + sta = (sta_t.unsqueeze(1) * sta_hw.unsqueeze(0)).reshape(T, T, H * W, H * W).transpose(1, 2) + return sta.reshape(T * H * W, T * H * W) + + def get_sparse_params(self, sample, device): + """ + Generate sparse attention parameters for the transformer based on sample dimensions. + + This method computes the sparse attention configuration needed for efficient video processing in the + transformer model. + + Args: + sample (torch.Tensor): Input sample tensor + device (torch.device): Device to place tensors on + + Returns: + Dict: Dictionary containing sparse attention parameters + """ + assert self.transformer.config.patch_size[0] == 1 + B, T, H, W, _ = sample.shape + T, H, W = ( + T // self.transformer.config.patch_size[0], + H // self.transformer.config.patch_size[1], + W // self.transformer.config.patch_size[2], + ) + if self.transformer.config.attention_type == "nabla": + sta_mask = self.fast_sta_nabla( + T, + H // 8, + W // 8, + self.transformer.config.attention_wT, + self.transformer.config.attention_wH, + self.transformer.config.attention_wW, + device=device, + ) + + sparse_params = { + "sta_mask": sta_mask.unsqueeze_(0).unsqueeze_(0), + "attention_type": self.transformer.config.attention_type, + "to_fractal": True, + "P": self.transformer.config.attention_P, + "wT": self.transformer.config.attention_wT, + "wW": self.transformer.config.attention_wW, + "wH": self.transformer.config.attention_wH, + "add_sta": self.transformer.config.attention_add_sta, + "visual_shape": (T, H, W), + "method": self.transformer.config.attention_method, + } + else: + sparse_params = None + + return sparse_params + + def _encode_prompt_qwen( + self, + prompt: Union[str, List[str]], + device: Optional[torch.device] = None, + max_sequence_length: int = 256, + dtype: Optional[torch.dtype] = None, + ): + """ + Encode prompt using Qwen2.5-VL text encoder. + + This method processes the input prompt through the Qwen2.5-VL model to generate text embeddings suitable for + video generation. + + Args: + prompt (Union[str, List[str]]): Input prompt or list of prompts + device (torch.device): Device to run encoding on + num_videos_per_prompt (int): Number of videos to generate per prompt + max_sequence_length (int): Maximum sequence length for tokenization + dtype (torch.dtype): Data type for embeddings + + Returns: + Tuple[torch.Tensor, torch.Tensor]: Text embeddings and cumulative sequence lengths + """ + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + + full_texts = [self.prompt_template.format(p) for p in prompt] + + inputs = self.tokenizer( + text=full_texts, + images=None, + videos=None, + max_length=max_sequence_length + self.prompt_template_encode_start_idx, + truncation=True, + return_tensors="pt", + padding=True, + ).to(device) + + embeds = self.text_encoder( + input_ids=inputs["input_ids"], + return_dict=True, + output_hidden_states=True, + )["hidden_states"][-1][:, self.prompt_template_encode_start_idx :] + + attention_mask = inputs["attention_mask"][:, self.prompt_template_encode_start_idx :] + cu_seqlens = torch.cumsum(attention_mask.sum(1), dim=0) + cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0).to(dtype=torch.int32) + + return embeds.to(dtype), cu_seqlens + + def _encode_prompt_clip( + self, + prompt: Union[str, List[str]], + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + """ + Encode prompt using CLIP text encoder. + + This method processes the input prompt through the CLIP model to generate pooled embeddings that capture + semantic information. + + Args: + prompt (Union[str, List[str]]): Input prompt or list of prompts + device (torch.device): Device to run encoding on + num_videos_per_prompt (int): Number of videos to generate per prompt + dtype (torch.dtype): Data type for embeddings + + Returns: + torch.Tensor: Pooled text embeddings from CLIP + """ + device = device or self._execution_device + dtype = dtype or self.text_encoder_2.dtype + + inputs = self.tokenizer_2( + prompt, + max_length=77, + truncation=True, + add_special_tokens=True, + padding="max_length", + return_tensors="pt", + ).to(device) + + pooled_embed = self.text_encoder_2(**inputs)["pooler_output"] + + return pooled_embed.to(dtype) + + def encode_prompt( + self, + prompt: Union[str, List[str]], + num_videos_per_prompt: int = 1, + max_sequence_length: int = 512, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + r""" + Encodes a single prompt (positive or negative) into text encoder hidden states. + + This method combines embeddings from both Qwen2.5-VL and CLIP text encoders to create comprehensive text + representations for video generation. + + Args: + prompt (`str` or `List[str]`): + Prompt to be encoded. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + Number of videos to generate per prompt. + max_sequence_length (`int`, *optional*, defaults to 512): + Maximum sequence length for text encoding. + device (`torch.device`, *optional*): + Torch device. + dtype (`torch.dtype`, *optional*): + Torch dtype. + + Returns: + Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + - Qwen text embeddings of shape (batch_size * num_videos_per_prompt, sequence_length, embedding_dim) + - CLIP pooled embeddings of shape (batch_size * num_videos_per_prompt, clip_embedding_dim) + - Cumulative sequence lengths (`cu_seqlens`) for Qwen embeddings of shape (batch_size * + num_videos_per_prompt + 1,) + """ + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + + batch_size = len(prompt) + + prompt = [prompt_clean(p) for p in prompt] + + # Encode with Qwen2.5-VL + prompt_embeds_qwen, prompt_cu_seqlens = self._encode_prompt_qwen( + prompt=prompt, + device=device, + max_sequence_length=max_sequence_length, + dtype=dtype, + ) + # prompt_embeds_qwen shape: [batch_size, seq_len, embed_dim] + + # Encode with CLIP + prompt_embeds_clip = self._encode_prompt_clip( + prompt=prompt, + device=device, + dtype=dtype, + ) + # prompt_embeds_clip shape: [batch_size, clip_embed_dim] + + # Repeat embeddings for num_videos_per_prompt + # Qwen embeddings: repeat sequence for each video, then reshape + prompt_embeds_qwen = prompt_embeds_qwen.repeat( + 1, num_videos_per_prompt, 1 + ) # [batch_size, seq_len * num_videos_per_prompt, embed_dim] + # Reshape to [batch_size * num_videos_per_prompt, seq_len, embed_dim] + prompt_embeds_qwen = prompt_embeds_qwen.view( + batch_size * num_videos_per_prompt, -1, prompt_embeds_qwen.shape[-1] + ) + + # CLIP embeddings: repeat for each video + prompt_embeds_clip = prompt_embeds_clip.repeat( + 1, num_videos_per_prompt, 1 + ) # [batch_size, num_videos_per_prompt, clip_embed_dim] + # Reshape to [batch_size * num_videos_per_prompt, clip_embed_dim] + prompt_embeds_clip = prompt_embeds_clip.view(batch_size * num_videos_per_prompt, -1) + + # Repeat cumulative sequence lengths for num_videos_per_prompt + # Original cu_seqlens: [0, len1, len1+len2, ...] + # Need to repeat the differences and reconstruct for repeated prompts + # Original differences (lengths) for each prompt in the batch + original_lengths = prompt_cu_seqlens.diff() # [len1, len2, ...] + # Repeat the lengths for num_videos_per_prompt + repeated_lengths = original_lengths.repeat_interleave( + num_videos_per_prompt + ) # [len1, len1, ..., len2, len2, ...] + # Reconstruct the cumulative lengths + repeated_cu_seqlens = torch.cat( + [torch.tensor([0], device=device, dtype=torch.int32), repeated_lengths.cumsum(0)] + ) + + return prompt_embeds_qwen, prompt_embeds_clip, repeated_cu_seqlens + + def check_inputs( + self, + prompt, + negative_prompt, + height, + width, + prompt_embeds_qwen=None, + prompt_embeds_clip=None, + negative_prompt_embeds_qwen=None, + negative_prompt_embeds_clip=None, + prompt_cu_seqlens=None, + negative_prompt_cu_seqlens=None, + callback_on_step_end_tensor_inputs=None, + ): + """ + Validate input parameters for the pipeline. + + Args: + prompt: Input prompt + negative_prompt: Negative prompt for guidance + height: Video height + width: Video width + prompt_embeds_qwen: Pre-computed Qwen prompt embeddings + prompt_embeds_clip: Pre-computed CLIP prompt embeddings + negative_prompt_embeds_qwen: Pre-computed Qwen negative prompt embeddings + negative_prompt_embeds_clip: Pre-computed CLIP negative prompt embeddings + prompt_cu_seqlens: Pre-computed cumulative sequence lengths for Qwen positive prompt + negative_prompt_cu_seqlens: Pre-computed cumulative sequence lengths for Qwen negative prompt + callback_on_step_end_tensor_inputs: Callback tensor inputs + + Raises: + ValueError: If inputs are invalid + """ + if height % 16 != 0 or width % 16 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.") + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + # Check for consistency within positive prompt embeddings and sequence lengths + if prompt_embeds_qwen is not None or prompt_embeds_clip is not None or prompt_cu_seqlens is not None: + if prompt_embeds_qwen is None or prompt_embeds_clip is None or prompt_cu_seqlens is None: + raise ValueError( + "If any of `prompt_embeds_qwen`, `prompt_embeds_clip`, or `prompt_cu_seqlens` is provided, " + "all three must be provided." + ) + + # Check for consistency within negative prompt embeddings and sequence lengths + if ( + negative_prompt_embeds_qwen is not None + or negative_prompt_embeds_clip is not None + or negative_prompt_cu_seqlens is not None + ): + if ( + negative_prompt_embeds_qwen is None + or negative_prompt_embeds_clip is None + or negative_prompt_cu_seqlens is None + ): + raise ValueError( + "If any of `negative_prompt_embeds_qwen`, `negative_prompt_embeds_clip`, or `negative_prompt_cu_seqlens` is provided, " + "all three must be provided." + ) + + # Check if prompt or embeddings are provided (either prompt or all required embedding components for positive) + if prompt is None and prompt_embeds_qwen is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds_qwen` (and corresponding `prompt_embeds_clip` and `prompt_cu_seqlens`). Cannot leave all undefined." + ) + + # Validate types for prompt and negative_prompt if provided + if prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + if negative_prompt is not None and ( + not isinstance(negative_prompt, str) and not isinstance(negative_prompt, list) + ): + raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}") + + def prepare_latents( + self, + batch_size: int, + num_channels_latents: int = 16, + height: int = 480, + width: int = 832, + num_frames: int = 81, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """ + Prepare initial latent variables for video generation. + + This method creates random noise latents or uses provided latents as starting point for the denoising process. + + Args: + batch_size (int): Number of videos to generate + num_channels_latents (int): Number of channels in latent space + height (int): Height of generated video + width (int): Width of generated video + num_frames (int): Number of frames in video + dtype (torch.dtype): Data type for latents + device (torch.device): Device to create latents on + generator (torch.Generator): Random number generator + latents (torch.Tensor): Pre-existing latents to use + + Returns: + torch.Tensor: Prepared latent tensor + """ + if latents is not None: + return latents.to(device=device, dtype=dtype) + + num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + shape = ( + batch_size, + num_latent_frames, + int(height) // self.vae_scale_factor_spatial, + int(width) // self.vae_scale_factor_spatial, + num_channels_latents, + ) + 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." + ) + + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + + if self.transformer.visual_cond: + # For visual conditioning, concatenate with zeros and mask + visual_cond = torch.zeros_like(latents) + visual_cond_mask = torch.zeros( + [ + batch_size, + num_latent_frames, + int(height) // self.vae_scale_factor_spatial, + int(width) // self.vae_scale_factor_spatial, + 1, + ], + dtype=latents.dtype, + device=latents.device, + ) + latents = torch.cat([latents, visual_cond, visual_cond_mask], dim=-1) + + return latents + + @property + def guidance_scale(self): + """Get the current guidance scale value.""" + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + """Check if classifier-free guidance is enabled.""" + return self._guidance_scale > 1.0 + + @property + def num_timesteps(self): + """Get the number of denoising timesteps.""" + return self._num_timesteps + + @property + def interrupt(self): + """Check if generation has been interrupted.""" + return self._interrupt + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: Union[str, List[str]] = None, + negative_prompt: Optional[Union[str, List[str]]] = None, + height: int = 512, + width: int = 768, + num_frames: int = 121, + num_inference_steps: int = 50, + guidance_scale: float = 5.0, + num_videos_per_prompt: Optional[int] = 1, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + prompt_embeds_qwen: Optional[torch.Tensor] = None, + prompt_embeds_clip: Optional[torch.Tensor] = None, + negative_prompt_embeds_qwen: Optional[torch.Tensor] = None, + negative_prompt_embeds_clip: Optional[torch.Tensor] = None, + prompt_cu_seqlens: Optional[torch.Tensor] = None, + negative_prompt_cu_seqlens: Optional[torch.Tensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + callback_on_step_end: Optional[ + Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] + ] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + max_sequence_length: int = 512, + **kwargs, + ): + r""" + The call function to the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the video generation. If not defined, pass `prompt_embeds` instead. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to avoid during video generation. If not defined, pass `negative_prompt_embeds` + instead. Ignored when not using guidance (`guidance_scale` < `1`). + height (`int`, defaults to `512`): + The height in pixels of the generated video. + width (`int`, defaults to `768`): + The width in pixels of the generated video. + num_frames (`int`, defaults to `25`): + The number of frames in the generated video. + num_inference_steps (`int`, defaults to `50`): + The number of denoising steps. + guidance_scale (`float`, defaults to `5.0`): + Guidance scale as defined in classifier-free guidance. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of videos to generate per prompt. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + A torch generator to make generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generated video. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`KandinskyPipelineOutput`]. + callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*): + A function that is called at the end of each denoising step. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. + max_sequence_length (`int`, defaults to `512`): + The maximum sequence length for text encoding. + + Examples: + + Returns: + [`~KandinskyPipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`KandinskyPipelineOutput`] is returned, otherwise a `tuple` is returned + where the first element is a list with the generated images. + """ + if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): + callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt=prompt, + negative_prompt=negative_prompt, + height=height, + width=width, + prompt_embeds_qwen=prompt_embeds_qwen, + prompt_embeds_clip=prompt_embeds_clip, + negative_prompt_embeds_qwen=negative_prompt_embeds_qwen, + negative_prompt_embeds_clip=negative_prompt_embeds_clip, + prompt_cu_seqlens=prompt_cu_seqlens, + negative_prompt_cu_seqlens=negative_prompt_cu_seqlens, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + ) + + if num_frames % self.vae_scale_factor_temporal != 1: + logger.warning( + f"`num_frames - 1` has to be divisible by {self.vae_scale_factor_temporal}. Rounding to the nearest number." + ) + num_frames = num_frames // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1 + num_frames = max(num_frames, 1) + + self._guidance_scale = guidance_scale + self._interrupt = False + + device = self._execution_device + dtype = self.transformer.dtype + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + prompt = [prompt] + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds_qwen.shape[0] + + # 3. Encode input prompt + if prompt_embeds_qwen is None: + prompt_embeds_qwen, prompt_embeds_clip, prompt_cu_seqlens = self.encode_prompt( + prompt=prompt, + max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, + ) + + if self.do_classifier_free_guidance: + if negative_prompt is None: + negative_prompt = "Static, 2D cartoon, cartoon, 2d animation, paintings, images, worst quality, low quality, ugly, deformed, walking backwards" + + if isinstance(negative_prompt, str): + negative_prompt = [negative_prompt] * len(prompt) if prompt is not None else [negative_prompt] + elif len(negative_prompt) != len(prompt): + raise ValueError( + f"`negative_prompt` must have same length as `prompt`. Got {len(negative_prompt)} vs {len(prompt)}." + ) + + if negative_prompt_embeds_qwen is None: + negative_prompt_embeds_qwen, negative_prompt_embeds_clip, negative_cu_seqlens = self.encode_prompt( + prompt=negative_prompt, + max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, + ) + + # 4. Prepare timesteps + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = self.scheduler.timesteps + + # 5. Prepare latent variables + num_channels_latents = self.transformer.config.in_visual_dim + latents = self.prepare_latents( + batch_size * num_videos_per_prompt, + num_channels_latents, + height, + width, + num_frames, + dtype, + device, + generator, + latents, + ) + + # 6. Prepare rope positions for positional encoding + num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + visual_rope_pos = [ + torch.arange(num_latent_frames, device=device), + torch.arange(height // self.vae_scale_factor_spatial // 2, device=device), + torch.arange(width // self.vae_scale_factor_spatial // 2, device=device), + ] + + text_rope_pos = torch.arange(prompt_cu_seqlens.diff().max().item(), device=device) + + negative_text_rope_pos = ( + torch.arange(negative_cu_seqlens.diff().max().item(), device=device) + if negative_cu_seqlens is not None + else None + ) + + # 7. Sparse Params for efficient attention + sparse_params = self.get_sparse_params(latents, device) + + # 8. Denoising loop + num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + self._num_timesteps = len(timesteps) + + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + timestep = t.unsqueeze(0).repeat(batch_size * num_videos_per_prompt) + + # Predict noise residual + pred_velocity = self.transformer( + hidden_states=latents.to(dtype), + encoder_hidden_states=prompt_embeds_qwen.to(dtype), + pooled_projections=prompt_embeds_clip.to(dtype), + timestep=timestep.to(dtype), + visual_rope_pos=visual_rope_pos, + text_rope_pos=text_rope_pos, + scale_factor=(1, 2, 2), + sparse_params=sparse_params, + return_dict=True, + ).sample + + if self.do_classifier_free_guidance and negative_prompt_embeds_qwen is not None: + uncond_pred_velocity = self.transformer( + hidden_states=latents.to(dtype), + encoder_hidden_states=negative_prompt_embeds_qwen.to(dtype), + pooled_projections=negative_prompt_embeds_clip.to(dtype), + timestep=timestep.to(dtype), + visual_rope_pos=visual_rope_pos, + text_rope_pos=negative_text_rope_pos, + scale_factor=(1, 2, 2), + sparse_params=sparse_params, + return_dict=True, + ).sample + + pred_velocity = uncond_pred_velocity + guidance_scale * (pred_velocity - uncond_pred_velocity) + # Compute previous sample using the scheduler + latents[:, :, :, :, :num_channels_latents] = self.scheduler.step( + pred_velocity, t, latents[:, :, :, :, :num_channels_latents], return_dict=False + )[0] + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds_qwen = callback_outputs.pop("prompt_embeds_qwen", prompt_embeds_qwen) + prompt_embeds_clip = callback_outputs.pop("prompt_embeds_clip", prompt_embeds_clip) + negative_prompt_embeds_qwen = callback_outputs.pop( + "negative_prompt_embeds_qwen", negative_prompt_embeds_qwen + ) + negative_prompt_embeds_clip = callback_outputs.pop( + "negative_prompt_embeds_clip", negative_prompt_embeds_clip + ) + + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + # 8. Post-processing - extract main latents + latents = latents[:, :, :, :, :num_channels_latents] + + # 9. Decode latents to video + if output_type != "latent": + latents = latents.to(self.vae.dtype) + # Reshape and normalize latents + video = latents.reshape( + batch_size, + num_videos_per_prompt, + (num_frames - 1) // self.vae_scale_factor_temporal + 1, + height // self.vae_scale_factor_spatial, + width // self.vae_scale_factor_spatial, + num_channels_latents, + ) + video = video.permute(0, 1, 5, 2, 3, 4) # [batch, num_videos, channels, frames, height, width] + video = video.reshape( + batch_size * num_videos_per_prompt, + num_channels_latents, + (num_frames - 1) // self.vae_scale_factor_temporal + 1, + height // self.vae_scale_factor_spatial, + width // self.vae_scale_factor_spatial, + ) + + # Normalize and decode through VAE + video = video / self.vae.config.scaling_factor + video = self.vae.decode(video).sample + video = self.video_processor.postprocess_video(video, output_type=output_type) + else: + video = latents + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (video,) + + return KandinskyPipelineOutput(frames=video) diff --git a/src/diffusers/pipelines/kandinsky5/pipeline_output.py b/src/diffusers/pipelines/kandinsky5/pipeline_output.py new file mode 100644 index 0000000000..ed77d42a9a --- /dev/null +++ b/src/diffusers/pipelines/kandinsky5/pipeline_output.py @@ -0,0 +1,20 @@ +from dataclasses import dataclass + +import torch + +from diffusers.utils import BaseOutput + + +@dataclass +class KandinskyPipelineOutput(BaseOutput): + r""" + Output class for Wan pipelines. + + Args: + frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]): + List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing + denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape + `(batch_size, num_frames, channels, height, width)`. + """ + + frames: torch.Tensor diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 6e7d227979..5d62709c28 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -918,6 +918,21 @@ class Kandinsky3UNet(metaclass=DummyObject): requires_backends(cls, ["torch"]) +class Kandinsky5Transformer3DModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class LatteTransformer3DModel(metaclass=DummyObject): _backends = ["torch"] diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 9ed6250452..3244ef12ef 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -1247,6 +1247,21 @@ class Kandinsky3Pipeline(metaclass=DummyObject): requires_backends(cls, ["torch", "transformers"]) +class Kandinsky5T2VPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class KandinskyCombinedPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] From 7853bfbed7310f7d46379e3cd0415002cf5e3eb8 Mon Sep 17 00:00:00 2001 From: dg845 <58458699+dg845@users.noreply.github.com> Date: Sun, 19 Oct 2025 18:41:58 -0700 Subject: [PATCH 075/123] Remove Qwen Image Redundant RoPE Cache (#12452) Refactor QwenEmbedRope to only use the LRU cache for RoPE caching --- .../transformers/transformer_qwenimage.py | 31 ++++++++++--------- 1 file changed, 17 insertions(+), 14 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_qwenimage.py b/src/diffusers/models/transformers/transformer_qwenimage.py index 05379270c1..c0fa031b9f 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage.py +++ b/src/diffusers/models/transformers/transformer_qwenimage.py @@ -180,7 +180,6 @@ class QwenEmbedRope(nn.Module): ], dim=1, ) - self.rope_cache = {} # DO NOT USING REGISTER BUFFER HERE, IT WILL CAUSE COMPLEX NUMBERS LOSE ITS IMAGINARY PART self.scale_rope = scale_rope @@ -195,10 +194,20 @@ class QwenEmbedRope(nn.Module): freqs = torch.polar(torch.ones_like(freqs), freqs) return freqs - def forward(self, video_fhw, txt_seq_lens, device): + def forward( + self, + video_fhw: Union[Tuple[int, int, int], List[Tuple[int, int, int]]], + txt_seq_lens: List[int], + device: torch.device, + ) -> Tuple[torch.Tensor, torch.Tensor]: """ - Args: video_fhw: [frame, height, width] a list of 3 integers representing the shape of the video Args: - txt_length: [bs] a list of 1 integers representing the length of the text + Args: + video_fhw (`Tuple[int, int, int]` or `List[Tuple[int, int, int]]`): + A list of 3 integers [frame, height, width] representing the shape of the video. + txt_seq_lens (`List[int]`): + A list of integers of length batch_size representing the length of each text prompt. + device: (`torch.device`): + The device on which to perform the RoPE computation. """ if self.pos_freqs.device != device: self.pos_freqs = self.pos_freqs.to(device) @@ -213,14 +222,8 @@ class QwenEmbedRope(nn.Module): max_vid_index = 0 for idx, fhw in enumerate(video_fhw): frame, height, width = fhw - rope_key = f"{idx}_{height}_{width}" - - if not torch.compiler.is_compiling(): - if rope_key not in self.rope_cache: - self.rope_cache[rope_key] = self._compute_video_freqs(frame, height, width, idx) - video_freq = self.rope_cache[rope_key] - else: - video_freq = self._compute_video_freqs(frame, height, width, idx) + # RoPE frequencies are cached via a lru_cache decorator on _compute_video_freqs + video_freq = self._compute_video_freqs(frame, height, width, idx) video_freq = video_freq.to(device) vid_freqs.append(video_freq) @@ -235,8 +238,8 @@ class QwenEmbedRope(nn.Module): return vid_freqs, txt_freqs - @functools.lru_cache(maxsize=None) - def _compute_video_freqs(self, frame, height, width, idx=0): + @functools.lru_cache(maxsize=128) + def _compute_video_freqs(self, frame: int, height: int, width: int, idx: int = 0) -> torch.Tensor: seq_lens = frame * height * width freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1) freqs_neg = self.neg_freqs.split([x // 2 for x in self.axes_dim], dim=1) From 48305755bfb1dd702f7f130aab767dfc2e1f67dc Mon Sep 17 00:00:00 2001 From: Dhruv Nair Date: Mon, 20 Oct 2025 22:32:23 +0530 Subject: [PATCH 076/123] Raise warning instead of error when imports are missing for custom code (#12513) update --- src/diffusers/utils/dynamic_modules_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/diffusers/utils/dynamic_modules_utils.py b/src/diffusers/utils/dynamic_modules_utils.py index 35bef51229..627b1e0604 100644 --- a/src/diffusers/utils/dynamic_modules_utils.py +++ b/src/diffusers/utils/dynamic_modules_utils.py @@ -151,8 +151,8 @@ def check_imports(filename): missing_packages.append(imp) if len(missing_packages) > 0: - raise ImportError( - "This modeling file requires the following packages that were not found in your environment: " + logger.warning( + "This modeling file might require the following packages that were not found in your environment: " f"{', '.join(missing_packages)}. Run `pip install {' '.join(missing_packages)}`" ) From decfa3c9e1c2ea55263102c6895d3385a0a7a1ad Mon Sep 17 00:00:00 2001 From: Fei Xie <69195059+FeiXie8@users.noreply.github.com> Date: Tue, 21 Oct 2025 09:45:37 +0800 Subject: [PATCH 077/123] =?UTF-8?q?Fix:=20Use=20incorrect=20temporary=20va?= =?UTF-8?q?riable=20key=20when=20replacing=20adapter=20name=E2=80=A6=20(#1?= =?UTF-8?q?2502)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fix: Use incorrect temporary variable key when replacing adapter name in state dict within load_lora_adapter function Co-authored-by: Sayak Paul --- src/diffusers/loaders/peft.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/loaders/peft.py b/src/diffusers/loaders/peft.py index 2381ccfef3..7d65b30659 100644 --- a/src/diffusers/loaders/peft.py +++ b/src/diffusers/loaders/peft.py @@ -293,7 +293,7 @@ class PeftAdapterMixin: # For hotswapping, we need the adapter name to be present in the state dict keys new_sd = {} for k, v in sd.items(): - if k.endswith("lora_A.weight") or key.endswith("lora_B.weight"): + if k.endswith("lora_A.weight") or k.endswith("lora_B.weight"): k = k[: -len(".weight")] + f".{adapter_name}.weight" elif k.endswith("lora_B.bias"): # lora_bias=True option k = k[: -len(".bias")] + f".{adapter_name}.bias" From 5b5fa49a899dd520d1ec9fd875e04d64028b5b19 Mon Sep 17 00:00:00 2001 From: Steven Liu <59462357+stevhliu@users.noreply.github.com> Date: Mon, 20 Oct 2025 21:48:54 -0700 Subject: [PATCH 078/123] [docs] Organize toctree by modality (#12514) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * reorganize * fix --------- Co-authored-by: Álvaro Somoza --- docs/source/en/_toctree.yml | 396 ++++++++++++++++++------------------ 1 file changed, 201 insertions(+), 195 deletions(-) diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 848e38079e..ffd7385d81 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -450,204 +450,210 @@ sections: - local: api/pipelines/overview title: Overview - - local: api/pipelines/allegro - title: Allegro - - local: api/pipelines/amused - title: aMUSEd - - local: api/pipelines/animatediff - title: AnimateDiff - - local: api/pipelines/attend_and_excite - title: Attend-and-Excite - - local: api/pipelines/audioldm - title: AudioLDM - - local: api/pipelines/audioldm2 - title: AudioLDM 2 - - local: api/pipelines/aura_flow - title: AuraFlow - local: api/pipelines/auto_pipeline title: AutoPipeline - - local: api/pipelines/blip_diffusion - title: BLIP-Diffusion - - local: api/pipelines/bria_3_2 - title: Bria 3.2 - - local: api/pipelines/chroma - title: Chroma - - local: api/pipelines/cogvideox - title: CogVideoX - - local: api/pipelines/cogview3 - title: CogView3 - - local: api/pipelines/cogview4 - title: CogView4 - - local: api/pipelines/consisid - title: ConsisID - - local: api/pipelines/consistency_models - title: Consistency Models - - local: api/pipelines/controlnet - title: ControlNet - - local: api/pipelines/controlnet_flux - title: ControlNet with Flux.1 - - local: api/pipelines/controlnet_hunyuandit - title: ControlNet with Hunyuan-DiT - - local: api/pipelines/controlnet_sd3 - title: ControlNet with Stable Diffusion 3 - - local: api/pipelines/controlnet_sdxl - title: ControlNet with Stable Diffusion XL - - local: api/pipelines/controlnet_sana - title: ControlNet-Sana - - local: api/pipelines/controlnetxs - title: ControlNet-XS - - local: api/pipelines/controlnetxs_sdxl - title: ControlNet-XS with Stable Diffusion XL - - local: api/pipelines/controlnet_union - title: ControlNetUnion - - local: api/pipelines/cosmos - title: Cosmos - - local: api/pipelines/dance_diffusion - title: Dance Diffusion - - local: api/pipelines/ddim - title: DDIM - - local: api/pipelines/ddpm - title: DDPM - - local: api/pipelines/deepfloyd_if - title: DeepFloyd IF - - local: api/pipelines/diffedit - title: DiffEdit - - local: api/pipelines/dit - title: DiT - - local: api/pipelines/easyanimate - title: EasyAnimate - - local: api/pipelines/flux - title: Flux - - local: api/pipelines/control_flux_inpaint - title: FluxControlInpaint - - local: api/pipelines/framepack - title: Framepack - - local: api/pipelines/hidream - title: HiDream-I1 - - local: api/pipelines/hunyuandit - title: Hunyuan-DiT - - local: api/pipelines/hunyuan_video - title: HunyuanVideo - - local: api/pipelines/i2vgenxl - title: I2VGen-XL - - local: api/pipelines/pix2pix - title: InstructPix2Pix - - local: api/pipelines/kandinsky - title: Kandinsky 2.1 - - local: api/pipelines/kandinsky_v22 - title: Kandinsky 2.2 - - local: api/pipelines/kandinsky3 - title: Kandinsky 3 - - local: api/pipelines/kolors - title: Kolors - - local: api/pipelines/latent_consistency_models - title: Latent Consistency Models - - local: api/pipelines/latent_diffusion - title: Latent Diffusion - - local: api/pipelines/latte - title: Latte - - local: api/pipelines/ledits_pp - title: LEDITS++ - - local: api/pipelines/ltx_video - title: LTXVideo - - local: api/pipelines/lumina2 - title: Lumina 2.0 - - local: api/pipelines/lumina - title: Lumina-T2X - - local: api/pipelines/marigold - title: Marigold - - local: api/pipelines/mochi - title: Mochi - - local: api/pipelines/panorama - title: MultiDiffusion - - local: api/pipelines/musicldm - title: MusicLDM - - local: api/pipelines/omnigen - title: OmniGen - - local: api/pipelines/pag - title: PAG - - local: api/pipelines/paint_by_example - title: Paint by Example - - local: api/pipelines/pia - title: Personalized Image Animator (PIA) - - local: api/pipelines/pixart - title: PixArt-α - - local: api/pipelines/pixart_sigma - title: PixArt-Σ - - local: api/pipelines/qwenimage - title: QwenImage - - local: api/pipelines/sana - title: Sana - - local: api/pipelines/sana_sprint - title: Sana Sprint - - local: api/pipelines/self_attention_guidance - title: Self-Attention Guidance - - local: api/pipelines/semantic_stable_diffusion - title: Semantic Guidance - - local: api/pipelines/shap_e - title: Shap-E - - local: api/pipelines/skyreels_v2 - title: SkyReels-V2 - - local: api/pipelines/stable_audio - title: Stable Audio - - local: api/pipelines/stable_cascade - title: Stable Cascade - - title: Stable Diffusion + - title: Image sections: - - local: api/pipelines/stable_diffusion/overview - title: Overview - - local: api/pipelines/stable_diffusion/depth2img - title: Depth-to-image - - local: api/pipelines/stable_diffusion/gligen - title: GLIGEN (Grounded Language-to-Image Generation) - - local: api/pipelines/stable_diffusion/image_variation - title: Image variation - - local: api/pipelines/stable_diffusion/img2img - title: Image-to-image + - local: api/pipelines/amused + title: aMUSEd + - local: api/pipelines/animatediff + title: AnimateDiff + - local: api/pipelines/attend_and_excite + title: Attend-and-Excite + - local: api/pipelines/aura_flow + title: AuraFlow + - local: api/pipelines/blip_diffusion + title: BLIP-Diffusion + - local: api/pipelines/bria_3_2 + title: Bria 3.2 + - local: api/pipelines/chroma + title: Chroma + - local: api/pipelines/cogview3 + title: CogView3 + - local: api/pipelines/cogview4 + title: CogView4 + - local: api/pipelines/consistency_models + title: Consistency Models + - local: api/pipelines/controlnet + title: ControlNet + - local: api/pipelines/controlnet_flux + title: ControlNet with Flux.1 + - local: api/pipelines/controlnet_hunyuandit + title: ControlNet with Hunyuan-DiT + - local: api/pipelines/controlnet_sd3 + title: ControlNet with Stable Diffusion 3 + - local: api/pipelines/controlnet_sdxl + title: ControlNet with Stable Diffusion XL + - local: api/pipelines/controlnet_sana + title: ControlNet-Sana + - local: api/pipelines/controlnetxs + title: ControlNet-XS + - local: api/pipelines/controlnetxs_sdxl + title: ControlNet-XS with Stable Diffusion XL + - local: api/pipelines/controlnet_union + title: ControlNetUnion + - local: api/pipelines/cosmos + title: Cosmos + - local: api/pipelines/ddim + title: DDIM + - local: api/pipelines/ddpm + title: DDPM + - local: api/pipelines/deepfloyd_if + title: DeepFloyd IF + - local: api/pipelines/diffedit + title: DiffEdit + - local: api/pipelines/dit + title: DiT + - local: api/pipelines/easyanimate + title: EasyAnimate + - local: api/pipelines/flux + title: Flux + - local: api/pipelines/control_flux_inpaint + title: FluxControlInpaint + - local: api/pipelines/hidream + title: HiDream-I1 + - local: api/pipelines/hunyuandit + title: Hunyuan-DiT + - local: api/pipelines/pix2pix + title: InstructPix2Pix + - local: api/pipelines/kandinsky + title: Kandinsky 2.1 + - local: api/pipelines/kandinsky_v22 + title: Kandinsky 2.2 + - local: api/pipelines/kandinsky3 + title: Kandinsky 3 + - local: api/pipelines/kolors + title: Kolors + - local: api/pipelines/latent_consistency_models + title: Latent Consistency Models + - local: api/pipelines/latent_diffusion + title: Latent Diffusion + - local: api/pipelines/ledits_pp + title: LEDITS++ + - local: api/pipelines/lumina2 + title: Lumina 2.0 + - local: api/pipelines/lumina + title: Lumina-T2X + - local: api/pipelines/marigold + title: Marigold + - local: api/pipelines/panorama + title: MultiDiffusion + - local: api/pipelines/omnigen + title: OmniGen + - local: api/pipelines/pag + title: PAG + - local: api/pipelines/paint_by_example + title: Paint by Example + - local: api/pipelines/pixart + title: PixArt-α + - local: api/pipelines/pixart_sigma + title: PixArt-Σ + - local: api/pipelines/qwenimage + title: QwenImage + - local: api/pipelines/sana + title: Sana + - local: api/pipelines/sana_sprint + title: Sana Sprint + - local: api/pipelines/self_attention_guidance + title: Self-Attention Guidance + - local: api/pipelines/semantic_stable_diffusion + title: Semantic Guidance + - local: api/pipelines/shap_e + title: Shap-E + - local: api/pipelines/stable_cascade + title: Stable Cascade + - title: Stable Diffusion + sections: + - local: api/pipelines/stable_diffusion/overview + title: Overview + - local: api/pipelines/stable_diffusion/depth2img + title: Depth-to-image + - local: api/pipelines/stable_diffusion/gligen + title: GLIGEN (Grounded Language-to-Image Generation) + - local: api/pipelines/stable_diffusion/image_variation + title: Image variation + - local: api/pipelines/stable_diffusion/img2img + title: Image-to-image + - local: api/pipelines/stable_diffusion/inpaint + title: Inpainting + - local: api/pipelines/stable_diffusion/k_diffusion + title: K-Diffusion + - local: api/pipelines/stable_diffusion/latent_upscale + title: Latent upscaler + - local: api/pipelines/stable_diffusion/ldm3d_diffusion + title: LDM3D Text-to-(RGB, Depth), Text-to-(RGB-pano, Depth-pano), LDM3D Upscaler + - local: api/pipelines/stable_diffusion/stable_diffusion_safe + title: Safe Stable Diffusion + - local: api/pipelines/stable_diffusion/sdxl_turbo + title: SDXL Turbo + - local: api/pipelines/stable_diffusion/stable_diffusion_2 + title: Stable Diffusion 2 + - local: api/pipelines/stable_diffusion/stable_diffusion_3 + title: Stable Diffusion 3 + - local: api/pipelines/stable_diffusion/stable_diffusion_xl + title: Stable Diffusion XL + - local: api/pipelines/stable_diffusion/upscale + title: Super-resolution + - local: api/pipelines/stable_diffusion/adapter + title: T2I-Adapter + - local: api/pipelines/stable_diffusion/text2img + title: Text-to-image + - local: api/pipelines/stable_unclip + title: Stable unCLIP + - local: api/pipelines/unclip + title: unCLIP + - local: api/pipelines/unidiffuser + title: UniDiffuser + - local: api/pipelines/value_guided_sampling + title: Value-guided sampling + - local: api/pipelines/visualcloze + title: VisualCloze + - local: api/pipelines/wuerstchen + title: Wuerstchen + - title: Video + sections: + - local: api/pipelines/allegro + title: Allegro + - local: api/pipelines/cogvideox + title: CogVideoX + - local: api/pipelines/consisid + title: ConsisID + - local: api/pipelines/framepack + title: Framepack + - local: api/pipelines/hunyuan_video + title: HunyuanVideo + - local: api/pipelines/i2vgenxl + title: I2VGen-XL + - local: api/pipelines/latte + title: Latte + - local: api/pipelines/ltx_video + title: LTXVideo + - local: api/pipelines/mochi + title: Mochi + - local: api/pipelines/pia + title: Personalized Image Animator (PIA) + - local: api/pipelines/skyreels_v2 + title: SkyReels-V2 - local: api/pipelines/stable_diffusion/svd - title: Image-to-video - - local: api/pipelines/stable_diffusion/inpaint - title: Inpainting - - local: api/pipelines/stable_diffusion/k_diffusion - title: K-Diffusion - - local: api/pipelines/stable_diffusion/latent_upscale - title: Latent upscaler - - local: api/pipelines/stable_diffusion/ldm3d_diffusion - title: LDM3D Text-to-(RGB, Depth), Text-to-(RGB-pano, Depth-pano), LDM3D Upscaler - - local: api/pipelines/stable_diffusion/stable_diffusion_safe - title: Safe Stable Diffusion - - local: api/pipelines/stable_diffusion/sdxl_turbo - title: SDXL Turbo - - local: api/pipelines/stable_diffusion/stable_diffusion_2 - title: Stable Diffusion 2 - - local: api/pipelines/stable_diffusion/stable_diffusion_3 - title: Stable Diffusion 3 - - local: api/pipelines/stable_diffusion/stable_diffusion_xl - title: Stable Diffusion XL - - local: api/pipelines/stable_diffusion/upscale - title: Super-resolution - - local: api/pipelines/stable_diffusion/adapter - title: T2I-Adapter - - local: api/pipelines/stable_diffusion/text2img - title: Text-to-image - - local: api/pipelines/stable_unclip - title: Stable unCLIP - - local: api/pipelines/text_to_video - title: Text-to-video - - local: api/pipelines/text_to_video_zero - title: Text2Video-Zero - - local: api/pipelines/unclip - title: unCLIP - - local: api/pipelines/unidiffuser - title: UniDiffuser - - local: api/pipelines/value_guided_sampling - title: Value-guided sampling - - local: api/pipelines/visualcloze - title: VisualCloze - - local: api/pipelines/wan - title: Wan - - local: api/pipelines/wuerstchen - title: Wuerstchen + title: Stable Video Diffusion + - local: api/pipelines/text_to_video + title: Text-to-video + - local: api/pipelines/text_to_video_zero + title: Text2Video-Zero + - local: api/pipelines/wan + title: Wan + - title: Audio + sections: + - local: api/pipelines/audioldm + title: AudioLDM + - local: api/pipelines/audioldm2 + title: AudioLDM 2 + - local: api/pipelines/dance_diffusion + title: Dance Diffusion + - local: api/pipelines/musicldm + title: MusicLDM + - local: api/pipelines/stable_audio + title: Stable Audio - title: Schedulers sections: - local: api/schedulers/overview @@ -735,4 +741,4 @@ - local: api/image_processor title: VAE Image Processor - local: api/video_processor - title: Video Processor + title: Video Processor \ No newline at end of file From b3e56e71fb7c73601851bb83e7583f113f563d26 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Tue, 21 Oct 2025 20:04:54 +0530 Subject: [PATCH 079/123] styling issues. (#12522) --- docs/source/en/_toctree.yml | 165 +++++++++++++++++------------------- 1 file changed, 77 insertions(+), 88 deletions(-) diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index ffd7385d81..85e996a1b7 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -1,5 +1,4 @@ -- title: Get started - sections: +- sections: - local: index title: Diffusers - local: installation @@ -8,9 +7,8 @@ title: Quickstart - local: stable_diffusion title: Basic performance - -- title: Pipelines - isExpanded: false + title: Get started +- isExpanded: false sections: - local: using-diffusers/loading title: DiffusionPipeline @@ -28,9 +26,8 @@ title: Model formats - local: using-diffusers/push_to_hub title: Sharing pipelines and models - -- title: Adapters - isExpanded: false + title: Pipelines +- isExpanded: false sections: - local: tutorials/using_peft_for_inference title: LoRA @@ -44,9 +41,8 @@ title: DreamBooth - local: using-diffusers/textual_inversion_inference title: Textual inversion - -- title: Inference - isExpanded: false + title: Adapters +- isExpanded: false sections: - local: using-diffusers/weighted_prompts title: Prompting @@ -56,9 +52,8 @@ title: Batch inference - local: training/distributed_inference title: Distributed inference - -- title: Inference optimization - isExpanded: false + title: Inference +- isExpanded: false sections: - local: optimization/fp16 title: Accelerate inference @@ -70,8 +65,7 @@ title: Reduce memory usage - local: optimization/speed-memory-optims title: Compiling and offloading quantized models - - title: Community optimizations - sections: + - sections: - local: optimization/pruna title: Pruna - local: optimization/xformers @@ -90,9 +84,9 @@ title: ParaAttention - local: using-diffusers/image_quality title: FreeU - -- title: Hybrid Inference - isExpanded: false + title: Community optimizations + title: Inference optimization +- isExpanded: false sections: - local: hybrid_inference/overview title: Overview @@ -102,9 +96,8 @@ title: VAE Encode - local: hybrid_inference/api_reference title: API Reference - -- title: Modular Diffusers - isExpanded: false + title: Hybrid Inference +- isExpanded: false sections: - local: modular_diffusers/overview title: Overview @@ -126,9 +119,8 @@ title: ComponentsManager - local: modular_diffusers/guiders title: Guiders - -- title: Training - isExpanded: false + title: Modular Diffusers +- isExpanded: false sections: - local: training/overview title: Overview @@ -138,8 +130,7 @@ title: Adapt a model to a new task - local: tutorials/basic_training title: Train a diffusion model - - title: Models - sections: + - sections: - local: training/unconditional_training title: Unconditional image generation - local: training/text2image @@ -158,8 +149,8 @@ title: InstructPix2Pix - local: training/cogvideox title: CogVideoX - - title: Methods - sections: + title: Models + - sections: - local: training/text_inversion title: Textual Inversion - local: training/dreambooth @@ -172,9 +163,9 @@ title: Latent Consistency Distillation - local: training/ddpo title: Reinforcement learning training with DDPO - -- title: Quantization - isExpanded: false + title: Methods + title: Training +- isExpanded: false sections: - local: quantization/overview title: Getting started @@ -188,9 +179,8 @@ title: quanto - local: quantization/modelopt title: NVIDIA ModelOpt - -- title: Model accelerators and hardware - isExpanded: false + title: Quantization +- isExpanded: false sections: - local: optimization/onnx title: ONNX @@ -204,9 +194,8 @@ title: Intel Gaudi - local: optimization/neuron title: AWS Neuron - -- title: Specific pipeline examples - isExpanded: false + title: Model accelerators and hardware +- isExpanded: false sections: - local: using-diffusers/consisid title: ConsisID @@ -232,12 +221,10 @@ title: Stable Video Diffusion - local: using-diffusers/marigold_usage title: Marigold Computer Vision - -- title: Resources - isExpanded: false + title: Specific pipeline examples +- isExpanded: false sections: - - title: Task recipes - sections: + - sections: - local: using-diffusers/unconditional_image_generation title: Unconditional image generation - local: using-diffusers/conditional_image_generation @@ -252,6 +239,7 @@ title: Video generation - local: using-diffusers/depth2img title: Depth-to-image + title: Task recipes - local: using-diffusers/write_own_pipeline title: Understanding pipelines, models and schedulers - local: community_projects @@ -266,12 +254,10 @@ title: Diffusers' Ethical Guidelines - local: conceptual/evaluation title: Evaluating Diffusion Models - -- title: API - isExpanded: false + title: Resources +- isExpanded: false sections: - - title: Main Classes - sections: + - sections: - local: api/configuration title: Configuration - local: api/logging @@ -282,8 +268,8 @@ title: Quantization - local: api/parallel title: Parallel inference - - title: Modular - sections: + title: Main Classes + - sections: - local: api/modular_diffusers/pipeline title: Pipeline - local: api/modular_diffusers/pipeline_blocks @@ -294,8 +280,8 @@ title: Components and configs - local: api/modular_diffusers/guiders title: Guiders - - title: Loaders - sections: + title: Modular + - sections: - local: api/loaders/ip_adapter title: IP-Adapter - local: api/loaders/lora @@ -310,14 +296,13 @@ title: SD3Transformer2D - local: api/loaders/peft title: PEFT - - title: Models - sections: + title: Loaders + - sections: - local: api/models/overview title: Overview - local: api/models/auto_model title: AutoModel - - title: ControlNets - sections: + - sections: - local: api/models/controlnet title: ControlNetModel - local: api/models/controlnet_union @@ -332,8 +317,8 @@ title: SD3ControlNetModel - local: api/models/controlnet_sparsectrl title: SparseControlNetModel - - title: Transformers - sections: + title: ControlNets + - sections: - local: api/models/allegro_transformer3d title: AllegroTransformer3DModel - local: api/models/aura_flow_transformer2d @@ -396,8 +381,8 @@ title: TransformerTemporalModel - local: api/models/wan_transformer_3d title: WanTransformer3DModel - - title: UNets - sections: + title: Transformers + - sections: - local: api/models/stable_cascade_unet title: StableCascadeUNet - local: api/models/unet @@ -412,8 +397,8 @@ title: UNetMotionModel - local: api/models/uvit2d title: UViT2DModel - - title: VAEs - sections: + title: UNets + - sections: - local: api/models/asymmetricautoencoderkl title: AsymmetricAutoencoderKL - local: api/models/autoencoder_dc @@ -446,14 +431,26 @@ title: Tiny AutoEncoder - local: api/models/vq title: VQModel - - title: Pipelines - sections: + title: VAEs + title: Models + - sections: - local: api/pipelines/overview title: Overview + - sections: + - local: api/pipelines/audioldm + title: AudioLDM + - local: api/pipelines/audioldm2 + title: AudioLDM 2 + - local: api/pipelines/dance_diffusion + title: Dance Diffusion + - local: api/pipelines/musicldm + title: MusicLDM + - local: api/pipelines/stable_audio + title: Stable Audio + title: Audio - local: api/pipelines/auto_pipeline title: AutoPipeline - - title: Image - sections: + - sections: - local: api/pipelines/amused title: aMUSEd - local: api/pipelines/animatediff @@ -562,8 +559,7 @@ title: Shap-E - local: api/pipelines/stable_cascade title: Stable Cascade - - title: Stable Diffusion - sections: + - sections: - local: api/pipelines/stable_diffusion/overview title: Overview - local: api/pipelines/stable_diffusion/depth2img @@ -581,7 +577,8 @@ - local: api/pipelines/stable_diffusion/latent_upscale title: Latent upscaler - local: api/pipelines/stable_diffusion/ldm3d_diffusion - title: LDM3D Text-to-(RGB, Depth), Text-to-(RGB-pano, Depth-pano), LDM3D Upscaler + title: LDM3D Text-to-(RGB, Depth), Text-to-(RGB-pano, Depth-pano), LDM3D + Upscaler - local: api/pipelines/stable_diffusion/stable_diffusion_safe title: Safe Stable Diffusion - local: api/pipelines/stable_diffusion/sdxl_turbo @@ -598,6 +595,7 @@ title: T2I-Adapter - local: api/pipelines/stable_diffusion/text2img title: Text-to-image + title: Stable Diffusion - local: api/pipelines/stable_unclip title: Stable unCLIP - local: api/pipelines/unclip @@ -610,8 +608,8 @@ title: VisualCloze - local: api/pipelines/wuerstchen title: Wuerstchen - - title: Video - sections: + title: Image + - sections: - local: api/pipelines/allegro title: Allegro - local: api/pipelines/cogvideox @@ -642,20 +640,9 @@ title: Text2Video-Zero - local: api/pipelines/wan title: Wan - - title: Audio - sections: - - local: api/pipelines/audioldm - title: AudioLDM - - local: api/pipelines/audioldm2 - title: AudioLDM 2 - - local: api/pipelines/dance_diffusion - title: Dance Diffusion - - local: api/pipelines/musicldm - title: MusicLDM - - local: api/pipelines/stable_audio - title: Stable Audio - - title: Schedulers - sections: + title: Video + title: Pipelines + - sections: - local: api/schedulers/overview title: Overview - local: api/schedulers/cm_stochastic_iterative @@ -724,8 +711,8 @@ title: UniPCMultistepScheduler - local: api/schedulers/vq_diffusion title: VQDiffusionScheduler - - title: Internal classes - sections: + title: Schedulers + - sections: - local: api/internal_classes_overview title: Overview - local: api/attnprocessor @@ -741,4 +728,6 @@ - local: api/image_processor title: VAE Image Processor - local: api/video_processor - title: Video Processor \ No newline at end of file + title: Video Processor + title: Internal classes + title: API From cefc2cf82dbdb5e4f725374420f0f6a91eb69048 Mon Sep 17 00:00:00 2001 From: David Bertoin Date: Tue, 21 Oct 2025 17:25:55 +0200 Subject: [PATCH 080/123] Add Photon model and pipeline support (#12456) * Add Photon model and pipeline support This commit adds support for the Photon image generation model: - PhotonTransformer2DModel: Core transformer architecture - PhotonPipeline: Text-to-image generation pipeline - Attention processor updates for Photon-specific attention mechanism - Conversion script for loading Photon checkpoints - Documentation and tests * just store the T5Gemma encoder * enhance_vae_properties if vae is provided only * remove autocast for text encoder forwad * BF16 example * conditioned CFG * remove enhance vae and use vae.config directly when possible * move PhotonAttnProcessor2_0 in transformer_photon * remove einops dependency and now inherits from AttentionMixin * unify the structure of the forward block * update doc * update doc * fix T5Gemma loading from hub * fix timestep shift * remove lora support from doc * Rename EmbedND for PhotoEmbedND * remove modulation dataclass * put _attn_forward and _ffn_forward logic in PhotonBlock's forward * renam LastLayer for FinalLayer * remove lora related code * rename vae_spatial_compression_ratio for vae_scale_factor * support prompt_embeds in call * move xattention conditionning out computation out of the denoising loop * add negative prompts * Use _import_structure for lazy loading * make quality + style * add pipeline test + corresponding fixes * utility function that determines the default resolution given the VAE * Refactor PhotonAttention to match Flux pattern * built-in RMSNorm * Revert accidental .gitignore change * parameter names match the standard diffusers conventions * renaming and remove unecessary attributes setting * Update docs/source/en/api/pipelines/photon.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * quantization example * added doc to toctree * Update docs/source/en/api/pipelines/photon.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/api/pipelines/photon.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/api/pipelines/photon.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * use dispatch_attention_fn for multiple attention backend support * naming changes * make fix copy * Update docs/source/en/api/pipelines/photon.md Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Add PhotonTransformer2DModel to TYPE_CHECKING imports * make fix-copies * Use Tuple instead of tuple Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * restrict the version of transformers Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update tests/pipelines/photon/test_pipeline_photon.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update tests/pipelines/photon/test_pipeline_photon.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * change | for Optional * fix nits. * use typing Dict --------- Co-authored-by: davidb Co-authored-by: David Briand Co-authored-by: davidb Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> Co-authored-by: sayakpaul --- docs/source/en/_toctree.yml | 2 + docs/source/en/api/pipelines/photon.md | 131 +++ scripts/convert_photon_to_diffusers.py | 345 ++++++++ src/diffusers/__init__.py | 4 + src/diffusers/models/__init__.py | 2 + src/diffusers/models/transformers/__init__.py | 1 + .../models/transformers/transformer_photon.py | 770 ++++++++++++++++++ src/diffusers/pipelines/__init__.py | 2 + src/diffusers/pipelines/photon/__init__.py | 63 ++ .../pipelines/photon/pipeline_output.py | 35 + .../pipelines/photon/pipeline_photon.py | 768 +++++++++++++++++ src/diffusers/utils/dummy_pt_objects.py | 15 + .../dummy_torch_and_transformers_objects.py | 15 + .../test_models_transformer_photon.py | 83 ++ tests/pipelines/photon/__init__.py | 0 .../pipelines/photon/test_pipeline_photon.py | 265 ++++++ 16 files changed, 2501 insertions(+) create mode 100644 docs/source/en/api/pipelines/photon.md create mode 100644 scripts/convert_photon_to_diffusers.py create mode 100644 src/diffusers/models/transformers/transformer_photon.py create mode 100644 src/diffusers/pipelines/photon/__init__.py create mode 100644 src/diffusers/pipelines/photon/pipeline_output.py create mode 100644 src/diffusers/pipelines/photon/pipeline_photon.py create mode 100644 tests/models/transformers/test_models_transformer_photon.py create mode 100644 tests/pipelines/photon/__init__.py create mode 100644 tests/pipelines/photon/test_pipeline_photon.py diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 85e996a1b7..3abe89437f 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -541,6 +541,8 @@ title: PAG - local: api/pipelines/paint_by_example title: Paint by Example + - local: api/pipelines/photon + title: Photon - local: api/pipelines/pixart title: PixArt-α - local: api/pipelines/pixart_sigma diff --git a/docs/source/en/api/pipelines/photon.md b/docs/source/en/api/pipelines/photon.md new file mode 100644 index 0000000000..293e05f0fd --- /dev/null +++ b/docs/source/en/api/pipelines/photon.md @@ -0,0 +1,131 @@ + + +# Photon + + +Photon generates high-quality images from text using a simplified MMDIT architecture where text tokens don't update through transformer blocks. It employs flow matching with discrete scheduling for efficient sampling and uses Google's T5Gemma-2B-2B-UL2 model for multi-language text encoding. The ~1.3B parameter transformer delivers fast inference without sacrificing quality. You can choose between Flux VAE (8x compression, 16 latent channels) for balanced quality and speed or DC-AE (32x compression, 32 latent channels) for latent compression and faster processing. + +## Available models + +Photon offers multiple variants with different VAE configurations, each optimized for specific resolutions. Base models excel with detailed prompts, capturing complex compositions and subtle details. Fine-tuned models trained on the [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) improve aesthetic quality, especially with simpler prompts. + + +| Model | Resolution | Fine-tuned | Distilled | Description | Suggested prompts | Suggested parameters | Recommended dtype | +|:-----:|:-----------------:|:----------:|:----------:|:----------:|:----------:|:----------:|:----------:| +| [`Photoroom/photon-256-t2i`](https://huggingface.co/Photoroom/photon-256-t2i)| 256 | No | No | Base model pre-trained at 256 with Flux VAE|Works best with detailed prompts in natural language|28 steps, cfg=5.0| `torch.bfloat16` | +| [`Photoroom/photon-256-t2i-sft`](https://huggingface.co/Photoroom/photon-256-t2i-sft)| 512 | Yes | No | Fine-tuned on the [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) dataset with Flux VAE | Can handle less detailed prompts|28 steps, cfg=5.0| `torch.bfloat16` | +| [`Photoroom/photon-512-t2i`](https://huggingface.co/Photoroom/photon-512-t2i)| 512 | No | No | Base model pre-trained at 512 with Flux VAE |Works best with detailed prompts in natural language|28 steps, cfg=5.0| `torch.bfloat16` | +| [`Photoroom/photon-512-t2i-sft`](https://huggingface.co/Photoroom/photon-512-t2i-sft)| 512 | Yes | No | Fine-tuned on the [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) dataset with Flux VAE | Can handle less detailed prompts in natural language|28 steps, cfg=5.0| `torch.bfloat16` | +| [`Photoroom/photon-512-t2i-sft-distilled`](https://huggingface.co/Photoroom/photon-512-t2i-sft-distilled)| 512 | Yes | Yes | 8-step distilled model from [`Photoroom/photon-512-t2i-sft`](https://huggingface.co/Photoroom/photon-512-t2i-sft) | Can handle less detailed prompts in natural language|8 steps, cfg=1.0| `torch.bfloat16` | +| [`Photoroom/photon-512-t2i-dc-ae`](https://huggingface.co/Photoroom/photon-512-t2i-dc-ae)| 512 | No | No | Base model pre-trained at 512 with [Deep Compression Autoencoder (DC-AE)](https://hanlab.mit.edu/projects/dc-ae)|Works best with detailed prompts in natural language|28 steps, cfg=5.0| `torch.bfloat16` | +| [`Photoroom/photon-512-t2i-dc-ae-sft`](https://huggingface.co/Photoroom/photon-512-t2i-dc-ae-sft)| 512 | Yes | No | Fine-tuned on the [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) dataset with [Deep Compression Autoencoder (DC-AE)](https://hanlab.mit.edu/projects/dc-ae) | Can handle less detailed prompts in natural language|28 steps, cfg=5.0| `torch.bfloat16` | +| [`Photoroom/photon-512-t2i-dc-ae-sft-distilled`](https://huggingface.co/Photoroom/photon-512-t2i-dc-ae-sft-distilled)| 512 | Yes | Yes | 8-step distilled model from [`Photoroom/photon-512-t2i-dc-ae-sft-distilled`](https://huggingface.co/Photoroom/photon-512-t2i-dc-ae-sft-distilled) | Can handle less detailed prompts in natural language|8 steps, cfg=1.0| `torch.bfloat16` |s + +Refer to [this](https://huggingface.co/collections/Photoroom/photon-models-68e66254c202ebfab99ad38e) collection for more information. + +## Loading the pipeline + +Load the pipeline with [`~DiffusionPipeline.from_pretrained`]. + +```py +from diffusers.pipelines.photon import PhotonPipeline + +# Load pipeline - VAE and text encoder will be loaded from HuggingFace +pipe = PhotonPipeline.from_pretrained("Photoroom/photon-512-t2i-sft", torch_dtype=torch.bfloat16) +pipe.to("cuda") + +prompt = "A front-facing portrait of a lion the golden savanna at sunset." +image = pipe(prompt, num_inference_steps=28, guidance_scale=5.0).images[0] +image.save("photon_output.png") +``` + +### Manual Component Loading + +Load components individually to customize the pipeline for instance to use quantized models. + +```py +import torch +from diffusers.pipelines.photon import PhotonPipeline +from diffusers.models import AutoencoderKL, AutoencoderDC +from diffusers.models.transformers.transformer_photon import PhotonTransformer2DModel +from diffusers.schedulers import FlowMatchEulerDiscreteScheduler +from transformers import T5GemmaModel, GemmaTokenizerFast +from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig +from transformers import BitsAndBytesConfig as BitsAndBytesConfig + +quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True) +# Load transformer +transformer = PhotonTransformer2DModel.from_pretrained( + "checkpoints/photon-512-t2i-sft", + subfolder="transformer", + quantization_config=quant_config, + torch_dtype=torch.bfloat16, +) + +# Load scheduler +scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + "checkpoints/photon-512-t2i-sft", subfolder="scheduler" +) + +# Load T5Gemma text encoder +t5gemma_model = T5GemmaModel.from_pretrained("google/t5gemma-2b-2b-ul2", + quantization_config=quant_config, + torch_dtype=torch.bfloat16) +text_encoder = t5gemma_model.encoder.to(dtype=torch.bfloat16) +tokenizer = GemmaTokenizerFast.from_pretrained("google/t5gemma-2b-2b-ul2") +tokenizer.model_max_length = 256 + +# Load VAE - choose either Flux VAE or DC-AE +# Flux VAE +vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", + subfolder="vae", + quantization_config=quant_config, + torch_dtype=torch.bfloat16) + +pipe = PhotonPipeline( + transformer=transformer, + scheduler=scheduler, + text_encoder=text_encoder, + tokenizer=tokenizer, + vae=vae +) +pipe.to("cuda") +``` + + +## Memory Optimization + +For memory-constrained environments: + +```py +import torch +from diffusers.pipelines.photon import PhotonPipeline + +pipe = PhotonPipeline.from_pretrained("Photoroom/photon-512-t2i-sft", torch_dtype=torch.bfloat16) +pipe.enable_model_cpu_offload() # Offload components to CPU when not in use + +# Or use sequential CPU offload for even lower memory +pipe.enable_sequential_cpu_offload() +``` + +## PhotonPipeline + +[[autodoc]] PhotonPipeline + - all + - __call__ + +## PhotonPipelineOutput + +[[autodoc]] pipelines.photon.pipeline_output.PhotonPipelineOutput diff --git a/scripts/convert_photon_to_diffusers.py b/scripts/convert_photon_to_diffusers.py new file mode 100644 index 0000000000..c66bc31418 --- /dev/null +++ b/scripts/convert_photon_to_diffusers.py @@ -0,0 +1,345 @@ +#!/usr/bin/env python3 +""" +Script to convert Photon checkpoint from original codebase to diffusers format. +""" + +import argparse +import json +import os +import sys +from dataclasses import asdict, dataclass +from typing import Dict, Tuple + +import torch +from safetensors.torch import save_file + +from diffusers.models.transformers.transformer_photon import PhotonTransformer2DModel +from diffusers.pipelines.photon import PhotonPipeline + + +DEFAULT_RESOLUTION = 512 + + +@dataclass(frozen=True) +class PhotonBase: + context_in_dim: int = 2304 + hidden_size: int = 1792 + mlp_ratio: float = 3.5 + num_heads: int = 28 + depth: int = 16 + axes_dim: Tuple[int, int] = (32, 32) + theta: int = 10_000 + time_factor: float = 1000.0 + time_max_period: int = 10_000 + + +@dataclass(frozen=True) +class PhotonFlux(PhotonBase): + in_channels: int = 16 + patch_size: int = 2 + + +@dataclass(frozen=True) +class PhotonDCAE(PhotonBase): + in_channels: int = 32 + patch_size: int = 1 + + +def build_config(vae_type: str) -> Tuple[dict, int]: + if vae_type == "flux": + cfg = PhotonFlux() + elif vae_type == "dc-ae": + cfg = PhotonDCAE() + else: + raise ValueError(f"Unsupported VAE type: {vae_type}. Use 'flux' or 'dc-ae'") + + config_dict = asdict(cfg) + config_dict["axes_dim"] = list(config_dict["axes_dim"]) # type: ignore[index] + return config_dict + + +def create_parameter_mapping(depth: int) -> dict: + """Create mapping from old parameter names to new diffusers names.""" + + # Key mappings for structural changes + mapping = {} + + # Map old structure (layers in PhotonBlock) to new structure (layers in PhotonAttention) + for i in range(depth): + # QKV projections moved to attention module + mapping[f"blocks.{i}.img_qkv_proj.weight"] = f"blocks.{i}.attention.img_qkv_proj.weight" + mapping[f"blocks.{i}.txt_kv_proj.weight"] = f"blocks.{i}.attention.txt_kv_proj.weight" + + # QK norm moved to attention module and renamed to match Attention's qk_norm structure + mapping[f"blocks.{i}.qk_norm.query_norm.scale"] = f"blocks.{i}.attention.norm_q.weight" + mapping[f"blocks.{i}.qk_norm.key_norm.scale"] = f"blocks.{i}.attention.norm_k.weight" + mapping[f"blocks.{i}.qk_norm.query_norm.weight"] = f"blocks.{i}.attention.norm_q.weight" + mapping[f"blocks.{i}.qk_norm.key_norm.weight"] = f"blocks.{i}.attention.norm_k.weight" + + # K norm for text tokens moved to attention module + mapping[f"blocks.{i}.k_norm.scale"] = f"blocks.{i}.attention.norm_added_k.weight" + mapping[f"blocks.{i}.k_norm.weight"] = f"blocks.{i}.attention.norm_added_k.weight" + + # Attention output projection + mapping[f"blocks.{i}.attn_out.weight"] = f"blocks.{i}.attention.to_out.0.weight" + + return mapping + + +def convert_checkpoint_parameters(old_state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]: + """Convert old checkpoint parameters to new diffusers format.""" + + print("Converting checkpoint parameters...") + + mapping = create_parameter_mapping(depth) + converted_state_dict = {} + + for key, value in old_state_dict.items(): + new_key = key + + # Apply specific mappings if needed + if key in mapping: + new_key = mapping[key] + print(f" Mapped: {key} -> {new_key}") + + converted_state_dict[new_key] = value + + print(f"✓ Converted {len(converted_state_dict)} parameters") + return converted_state_dict + + +def create_transformer_from_checkpoint(checkpoint_path: str, config: dict) -> PhotonTransformer2DModel: + """Create and load PhotonTransformer2DModel from old checkpoint.""" + + print(f"Loading checkpoint from: {checkpoint_path}") + + # Load old checkpoint + if not os.path.exists(checkpoint_path): + raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}") + + old_checkpoint = torch.load(checkpoint_path, map_location="cpu") + + # Handle different checkpoint formats + if isinstance(old_checkpoint, dict): + if "model" in old_checkpoint: + state_dict = old_checkpoint["model"] + elif "state_dict" in old_checkpoint: + state_dict = old_checkpoint["state_dict"] + else: + state_dict = old_checkpoint + else: + state_dict = old_checkpoint + + print(f"✓ Loaded checkpoint with {len(state_dict)} parameters") + + # Convert parameter names if needed + model_depth = int(config.get("depth", 16)) + converted_state_dict = convert_checkpoint_parameters(state_dict, depth=model_depth) + + # Create transformer with config + print("Creating PhotonTransformer2DModel...") + transformer = PhotonTransformer2DModel(**config) + + # Load state dict + print("Loading converted parameters...") + missing_keys, unexpected_keys = transformer.load_state_dict(converted_state_dict, strict=False) + + if missing_keys: + print(f"⚠ Missing keys: {missing_keys}") + if unexpected_keys: + print(f"⚠ Unexpected keys: {unexpected_keys}") + + if not missing_keys and not unexpected_keys: + print("✓ All parameters loaded successfully!") + + return transformer + + +def create_scheduler_config(output_path: str, shift: float): + """Create FlowMatchEulerDiscreteScheduler config.""" + + scheduler_config = {"_class_name": "FlowMatchEulerDiscreteScheduler", "num_train_timesteps": 1000, "shift": shift} + + scheduler_path = os.path.join(output_path, "scheduler") + os.makedirs(scheduler_path, exist_ok=True) + + with open(os.path.join(scheduler_path, "scheduler_config.json"), "w") as f: + json.dump(scheduler_config, f, indent=2) + + print("✓ Created scheduler config") + + +def download_and_save_vae(vae_type: str, output_path: str): + """Download and save VAE to local directory.""" + from diffusers import AutoencoderDC, AutoencoderKL + + vae_path = os.path.join(output_path, "vae") + os.makedirs(vae_path, exist_ok=True) + + if vae_type == "flux": + print("Downloading FLUX VAE from black-forest-labs/FLUX.1-dev...") + vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae") + else: # dc-ae + print("Downloading DC-AE VAE from mit-han-lab/dc-ae-f32c32-sana-1.1-diffusers...") + vae = AutoencoderDC.from_pretrained("mit-han-lab/dc-ae-f32c32-sana-1.1-diffusers") + + vae.save_pretrained(vae_path) + print(f"✓ Saved VAE to {vae_path}") + + +def download_and_save_text_encoder(output_path: str): + """Download and save T5Gemma text encoder and tokenizer.""" + from transformers import GemmaTokenizerFast + from transformers.models.t5gemma.modeling_t5gemma import T5GemmaModel + + text_encoder_path = os.path.join(output_path, "text_encoder") + tokenizer_path = os.path.join(output_path, "tokenizer") + os.makedirs(text_encoder_path, exist_ok=True) + os.makedirs(tokenizer_path, exist_ok=True) + + print("Downloading T5Gemma model from google/t5gemma-2b-2b-ul2...") + t5gemma_model = T5GemmaModel.from_pretrained("google/t5gemma-2b-2b-ul2") + + # Extract and save only the encoder + t5gemma_encoder = t5gemma_model.encoder + t5gemma_encoder.save_pretrained(text_encoder_path) + print(f"✓ Saved T5GemmaEncoder to {text_encoder_path}") + + print("Downloading tokenizer from google/t5gemma-2b-2b-ul2...") + tokenizer = GemmaTokenizerFast.from_pretrained("google/t5gemma-2b-2b-ul2") + tokenizer.model_max_length = 256 + tokenizer.save_pretrained(tokenizer_path) + print(f"✓ Saved tokenizer to {tokenizer_path}") + + +def create_model_index(vae_type: str, default_image_size: int, output_path: str): + """Create model_index.json for the pipeline.""" + + if vae_type == "flux": + vae_class = "AutoencoderKL" + else: # dc-ae + vae_class = "AutoencoderDC" + + model_index = { + "_class_name": "PhotonPipeline", + "_diffusers_version": "0.31.0.dev0", + "_name_or_path": os.path.basename(output_path), + "default_sample_size": default_image_size, + "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"], + "text_encoder": ["photon", "T5GemmaEncoder"], + "tokenizer": ["transformers", "GemmaTokenizerFast"], + "transformer": ["diffusers", "PhotonTransformer2DModel"], + "vae": ["diffusers", vae_class], + } + + model_index_path = os.path.join(output_path, "model_index.json") + with open(model_index_path, "w") as f: + json.dump(model_index, f, indent=2) + + +def main(args): + # Validate inputs + if not os.path.exists(args.checkpoint_path): + raise FileNotFoundError(f"Checkpoint not found: {args.checkpoint_path}") + + config = build_config(args.vae_type) + + # Create output directory + os.makedirs(args.output_path, exist_ok=True) + print(f"✓ Output directory: {args.output_path}") + + # Create transformer from checkpoint + transformer = create_transformer_from_checkpoint(args.checkpoint_path, config) + + # Save transformer + transformer_path = os.path.join(args.output_path, "transformer") + os.makedirs(transformer_path, exist_ok=True) + + # Save config + with open(os.path.join(transformer_path, "config.json"), "w") as f: + json.dump(config, f, indent=2) + + # Save model weights as safetensors + state_dict = transformer.state_dict() + save_file(state_dict, os.path.join(transformer_path, "diffusion_pytorch_model.safetensors")) + print(f"✓ Saved transformer to {transformer_path}") + + # Create scheduler config + create_scheduler_config(args.output_path, args.shift) + + download_and_save_vae(args.vae_type, args.output_path) + download_and_save_text_encoder(args.output_path) + + # Create model_index.json + create_model_index(args.vae_type, args.resolution, args.output_path) + + # Verify the pipeline can be loaded + try: + pipeline = PhotonPipeline.from_pretrained(args.output_path) + print("Pipeline loaded successfully!") + print(f"Transformer: {type(pipeline.transformer).__name__}") + print(f"VAE: {type(pipeline.vae).__name__}") + print(f"Text Encoder: {type(pipeline.text_encoder).__name__}") + print(f"Scheduler: {type(pipeline.scheduler).__name__}") + + # Display model info + num_params = sum(p.numel() for p in pipeline.transformer.parameters()) + print(f"✓ Transformer parameters: {num_params:,}") + + except Exception as e: + print(f"Pipeline verification failed: {e}") + return False + + print("Conversion completed successfully!") + print(f"Converted pipeline saved to: {args.output_path}") + print(f"VAE type: {args.vae_type}") + + return True + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Convert Photon checkpoint to diffusers format") + + parser.add_argument( + "--checkpoint_path", type=str, required=True, help="Path to the original Photon checkpoint (.pth file )" + ) + + parser.add_argument( + "--output_path", type=str, required=True, help="Output directory for the converted diffusers pipeline" + ) + + parser.add_argument( + "--vae_type", + type=str, + choices=["flux", "dc-ae"], + required=True, + help="VAE type to use: 'flux' for AutoencoderKL (16 channels) or 'dc-ae' for AutoencoderDC (32 channels)", + ) + + parser.add_argument( + "--resolution", + type=int, + choices=[256, 512, 1024], + default=DEFAULT_RESOLUTION, + help="Target resolution for the model (256, 512, or 1024). Affects the transformer's sample_size.", + ) + + parser.add_argument( + "--shift", + type=float, + default=3.0, + help="Shift for the scheduler", + ) + + args = parser.parse_args() + + try: + success = main(args) + if not success: + sys.exit(1) + except Exception as e: + print(f"Conversion failed: {e}") + import traceback + + traceback.print_exc() + sys.exit(1) diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index aa500b1494..b7086d2e0c 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -232,6 +232,7 @@ else: "MultiControlNetModel", "OmniGenTransformer2DModel", "ParallelConfig", + "PhotonTransformer2DModel", "PixArtTransformer2DModel", "PriorTransformer", "QwenImageControlNetModel", @@ -515,6 +516,7 @@ else: "MusicLDMPipeline", "OmniGenPipeline", "PaintByExamplePipeline", + "PhotonPipeline", "PIAPipeline", "PixArtAlphaPipeline", "PixArtSigmaPAGPipeline", @@ -926,6 +928,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: MultiControlNetModel, OmniGenTransformer2DModel, ParallelConfig, + PhotonTransformer2DModel, PixArtTransformer2DModel, PriorTransformer, QwenImageControlNetModel, @@ -1179,6 +1182,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: MusicLDMPipeline, OmniGenPipeline, PaintByExamplePipeline, + PhotonPipeline, PIAPipeline, PixArtAlphaPipeline, PixArtSigmaPAGPipeline, diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index 8d029bf5d3..2151e602b2 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -96,6 +96,7 @@ if is_torch_available(): _import_structure["transformers.transformer_lumina2"] = ["Lumina2Transformer2DModel"] _import_structure["transformers.transformer_mochi"] = ["MochiTransformer3DModel"] _import_structure["transformers.transformer_omnigen"] = ["OmniGenTransformer2DModel"] + _import_structure["transformers.transformer_photon"] = ["PhotonTransformer2DModel"] _import_structure["transformers.transformer_qwenimage"] = ["QwenImageTransformer2DModel"] _import_structure["transformers.transformer_sd3"] = ["SD3Transformer2DModel"] _import_structure["transformers.transformer_skyreels_v2"] = ["SkyReelsV2Transformer3DModel"] @@ -190,6 +191,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: LuminaNextDiT2DModel, MochiTransformer3DModel, OmniGenTransformer2DModel, + PhotonTransformer2DModel, PixArtTransformer2DModel, PriorTransformer, QwenImageTransformer2DModel, diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index 6b80ea6c82..ab5311518b 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -32,6 +32,7 @@ if is_torch_available(): from .transformer_lumina2 import Lumina2Transformer2DModel from .transformer_mochi import MochiTransformer3DModel from .transformer_omnigen import OmniGenTransformer2DModel + from .transformer_photon import PhotonTransformer2DModel from .transformer_qwenimage import QwenImageTransformer2DModel from .transformer_sd3 import SD3Transformer2DModel from .transformer_skyreels_v2 import SkyReelsV2Transformer3DModel diff --git a/src/diffusers/models/transformers/transformer_photon.py b/src/diffusers/models/transformers/transformer_photon.py new file mode 100644 index 0000000000..6314020c1c --- /dev/null +++ b/src/diffusers/models/transformers/transformer_photon.py @@ -0,0 +1,770 @@ +# Copyright 2025 The Photoroom and The HuggingFace Teams. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +from torch import nn +from torch.nn.functional import fold, unfold + +from ...configuration_utils import ConfigMixin, register_to_config +from ...utils import logging +from ..attention import AttentionMixin, AttentionModuleMixin +from ..attention_dispatch import dispatch_attention_fn +from ..embeddings import get_timestep_embedding +from ..modeling_outputs import Transformer2DModelOutput +from ..modeling_utils import ModelMixin +from ..normalization import RMSNorm + + +logger = logging.get_logger(__name__) + + +def get_image_ids(batch_size: int, height: int, width: int, patch_size: int, device: torch.device) -> torch.Tensor: + r""" + Generates 2D patch coordinate indices for a batch of images. + + Args: + batch_size (`int`): + Number of images in the batch. + height (`int`): + Height of the input images (in pixels). + width (`int`): + Width of the input images (in pixels). + patch_size (`int`): + Size of the square patches that the image is divided into. + device (`torch.device`): + The device on which to create the tensor. + + Returns: + `torch.Tensor`: + Tensor of shape `(batch_size, num_patches, 2)` containing the (row, col) coordinates of each patch in the + image grid. + """ + + img_ids = torch.zeros(height // patch_size, width // patch_size, 2, device=device) + img_ids[..., 0] = torch.arange(height // patch_size, device=device)[:, None] + img_ids[..., 1] = torch.arange(width // patch_size, device=device)[None, :] + return img_ids.reshape((height // patch_size) * (width // patch_size), 2).unsqueeze(0).repeat(batch_size, 1, 1) + + +def apply_rope(xq: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor: + r""" + Applies rotary positional embeddings (RoPE) to a query tensor. + + Args: + xq (`torch.Tensor`): + Input tensor of shape `(..., dim)` representing the queries. + freqs_cis (`torch.Tensor`): + Precomputed rotary frequency components of shape `(..., dim/2, 2)` containing cosine and sine pairs. + + Returns: + `torch.Tensor`: + Tensor of the same shape as `xq` with rotary embeddings applied. + """ + xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2) + # Ensure freqs_cis is on the same device as queries to avoid device mismatches with offloading + freqs_cis = freqs_cis.to(device=xq.device, dtype=xq_.dtype) + xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1] + return xq_out.reshape(*xq.shape).type_as(xq) + + +class PhotonAttnProcessor2_0: + r""" + Processor for implementing Photon-style attention with multi-source tokens and RoPE. Supports multiple attention + backends (Flash Attention, Sage Attention, etc.) via dispatch_attention_fn. + """ + + _attention_backend = None + _parallel_config = None + + def __init__(self): + if not hasattr(torch.nn.functional, "scaled_dot_product_attention"): + raise ImportError("PhotonAttnProcessor2_0 requires PyTorch 2.0, please upgrade PyTorch to 2.0.") + + def __call__( + self, + attn: "PhotonAttention", + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + image_rotary_emb: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + """ + Apply Photon attention using PhotonAttention module. + + Args: + attn: PhotonAttention module containing projection layers + hidden_states: Image tokens [B, L_img, D] + encoder_hidden_states: Text tokens [B, L_txt, D] + attention_mask: Boolean mask for text tokens [B, L_txt] + image_rotary_emb: Rotary positional embeddings [B, 1, L_img, head_dim//2, 2, 2] + """ + + if encoder_hidden_states is None: + raise ValueError("PhotonAttnProcessor2_0 requires 'encoder_hidden_states' containing text tokens.") + + # Project image tokens to Q, K, V + img_qkv = attn.img_qkv_proj(hidden_states) + B, L_img, _ = img_qkv.shape + img_qkv = img_qkv.reshape(B, L_img, 3, attn.heads, attn.head_dim) + img_qkv = img_qkv.permute(2, 0, 3, 1, 4) # [3, B, H, L_img, D] + img_q, img_k, img_v = img_qkv[0], img_qkv[1], img_qkv[2] + + # Apply QK normalization to image tokens + img_q = attn.norm_q(img_q) + img_k = attn.norm_k(img_k) + + # Project text tokens to K, V + txt_kv = attn.txt_kv_proj(encoder_hidden_states) + B, L_txt, _ = txt_kv.shape + txt_kv = txt_kv.reshape(B, L_txt, 2, attn.heads, attn.head_dim) + txt_kv = txt_kv.permute(2, 0, 3, 1, 4) # [2, B, H, L_txt, D] + txt_k, txt_v = txt_kv[0], txt_kv[1] + + # Apply K normalization to text tokens + txt_k = attn.norm_added_k(txt_k) + + # Apply RoPE to image queries and keys + if image_rotary_emb is not None: + img_q = apply_rope(img_q, image_rotary_emb) + img_k = apply_rope(img_k, image_rotary_emb) + + # Concatenate text and image keys/values + k = torch.cat((txt_k, img_k), dim=2) # [B, H, L_txt + L_img, D] + v = torch.cat((txt_v, img_v), dim=2) # [B, H, L_txt + L_img, D] + + # Build attention mask if provided + attn_mask_tensor = None + if attention_mask is not None: + bs, _, l_img, _ = img_q.shape + l_txt = txt_k.shape[2] + + if attention_mask.dim() != 2: + raise ValueError(f"Unsupported attention_mask shape: {attention_mask.shape}") + if attention_mask.shape[-1] != l_txt: + raise ValueError(f"attention_mask last dim {attention_mask.shape[-1]} must equal text length {l_txt}") + + device = img_q.device + ones_img = torch.ones((bs, l_img), dtype=torch.bool, device=device) + attention_mask = attention_mask.to(device=device, dtype=torch.bool) + joint_mask = torch.cat([attention_mask, ones_img], dim=-1) + attn_mask_tensor = joint_mask[:, None, None, :].expand(-1, attn.heads, l_img, -1) + + # Apply attention using dispatch_attention_fn for backend support + # Reshape to match dispatch_attention_fn expectations: [B, L, H, D] + query = img_q.transpose(1, 2) # [B, L_img, H, D] + key = k.transpose(1, 2) # [B, L_txt + L_img, H, D] + value = v.transpose(1, 2) # [B, L_txt + L_img, H, D] + + attn_output = dispatch_attention_fn( + query, + key, + value, + attn_mask=attn_mask_tensor, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + + # Reshape from [B, L_img, H, D] to [B, L_img, H*D] + batch_size, seq_len, num_heads, head_dim = attn_output.shape + attn_output = attn_output.reshape(batch_size, seq_len, num_heads * head_dim) + + # Apply output projection + attn_output = attn.to_out[0](attn_output) + if len(attn.to_out) > 1: + attn_output = attn.to_out[1](attn_output) # dropout if present + + return attn_output + + +class PhotonAttention(nn.Module, AttentionModuleMixin): + r""" + Photon-style attention module that handles multi-source tokens and RoPE. Similar to FluxAttention but adapted for + Photon's architecture. + """ + + _default_processor_cls = PhotonAttnProcessor2_0 + _available_processors = [PhotonAttnProcessor2_0] + + def __init__( + self, + query_dim: int, + heads: int = 8, + dim_head: int = 64, + bias: bool = False, + out_bias: bool = False, + eps: float = 1e-6, + processor=None, + ): + super().__init__() + + self.heads = heads + self.head_dim = dim_head + self.inner_dim = dim_head * heads + self.query_dim = query_dim + + self.img_qkv_proj = nn.Linear(query_dim, query_dim * 3, bias=bias) + + self.norm_q = RMSNorm(self.head_dim, eps=eps, elementwise_affine=True) + self.norm_k = RMSNorm(self.head_dim, eps=eps, elementwise_affine=True) + + self.txt_kv_proj = nn.Linear(query_dim, query_dim * 2, bias=bias) + self.norm_added_k = RMSNorm(self.head_dim, eps=eps, elementwise_affine=True) + + self.to_out = nn.ModuleList([]) + self.to_out.append(nn.Linear(self.inner_dim, query_dim, bias=out_bias)) + self.to_out.append(nn.Dropout(0.0)) + + if processor is None: + processor = self._default_processor_cls() + self.set_processor(processor) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + image_rotary_emb: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + return self.processor( + self, + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + image_rotary_emb=image_rotary_emb, + **kwargs, + ) + + +# inspired from https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/layers.py +class PhotonEmbedND(nn.Module): + r""" + N-dimensional rotary positional embedding. + + This module creates rotary embeddings (RoPE) across multiple axes, where each axis can have its own embedding + dimension. The embeddings are combined and returned as a single tensor + + Args: + dim (int): + Base embedding dimension (must be even). + theta (int): + Scaling factor that controls the frequency spectrum of the rotary embeddings. + axes_dim (list[int]): + List of embedding dimensions for each axis (each must be even). + """ + + def __init__(self, dim: int, theta: int, axes_dim: List[int]): + super().__init__() + self.dim = dim + self.theta = theta + self.axes_dim = axes_dim + + def rope(self, pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor: + assert dim % 2 == 0 + scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim + omega = 1.0 / (theta**scale) + out = pos.unsqueeze(-1) * omega.unsqueeze(0) + out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1) + # Native PyTorch equivalent of: Rearrange("b n d (i j) -> b n d i j", i=2, j=2) + # out shape: (b, n, d, 4) -> reshape to (b, n, d, 2, 2) + out = out.reshape(*out.shape[:-1], 2, 2) + return out.float() + + def forward(self, ids: torch.Tensor) -> torch.Tensor: + n_axes = ids.shape[-1] + emb = torch.cat( + [self.rope(ids[:, :, i], self.axes_dim[i], self.theta) for i in range(n_axes)], + dim=-3, + ) + return emb.unsqueeze(1) + + +class MLPEmbedder(nn.Module): + r""" + A simple 2-layer MLP used for embedding inputs. + + Args: + in_dim (`int`): + Dimensionality of the input features. + hidden_dim (`int`): + Dimensionality of the hidden and output embedding space. + + Returns: + `torch.Tensor`: + Tensor of shape `(..., hidden_dim)` containing the embedded representations. + """ + + def __init__(self, in_dim: int, hidden_dim: int): + super().__init__() + self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True) + self.silu = nn.SiLU() + self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.out_layer(self.silu(self.in_layer(x))) + + +class Modulation(nn.Module): + r""" + Modulation network that generates scale, shift, and gating parameters. + + Given an input vector, the module projects it through a linear layer to produce six chunks, which are grouped into + two tuples `(shift, scale, gate)`. + + Args: + dim (`int`): + Dimensionality of the input vector. The output will have `6 * dim` features internally. + + Returns: + ((`torch.Tensor`, `torch.Tensor`, `torch.Tensor`), (`torch.Tensor`, `torch.Tensor`, `torch.Tensor`)): + Two tuples `(shift, scale, gate)`. + """ + + def __init__(self, dim: int): + super().__init__() + self.lin = nn.Linear(dim, 6 * dim, bias=True) + nn.init.constant_(self.lin.weight, 0) + nn.init.constant_(self.lin.bias, 0) + + def forward( + self, vec: torch.Tensor + ) -> Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(6, dim=-1) + return tuple(out[:3]), tuple(out[3:]) + + +class PhotonBlock(nn.Module): + r""" + Multimodal transformer block with text–image cross-attention, modulation, and MLP. + + Args: + hidden_size (`int`): + Dimension of the hidden representations. + num_heads (`int`): + Number of attention heads. + mlp_ratio (`float`, *optional*, defaults to 4.0): + Expansion ratio for the hidden dimension inside the MLP. + qk_scale (`float`, *optional*): + Scale factor for queries and keys. If not provided, defaults to ``head_dim**-0.5``. + + Attributes: + img_pre_norm (`nn.LayerNorm`): + Pre-normalization applied to image tokens before attention. + attention (`PhotonAttention`): + Multi-head attention module with built-in QKV projections and normalizations for cross-attention between + image and text tokens. + post_attention_layernorm (`nn.LayerNorm`): + Normalization applied after attention. + gate_proj / up_proj / down_proj (`nn.Linear`): + Feedforward layers forming the gated MLP. + mlp_act (`nn.GELU`): + Nonlinear activation used in the MLP. + modulation (`Modulation`): + Produces scale/shift/gating parameters for modulated layers. + + Methods: + The forward method performs cross-attention and the MLP with modulation. + """ + + def __init__( + self, + hidden_size: int, + num_heads: int, + mlp_ratio: float = 4.0, + qk_scale: Optional[float] = None, + ): + super().__init__() + + self.hidden_dim = hidden_size + self.num_heads = num_heads + self.head_dim = hidden_size // num_heads + self.scale = qk_scale or self.head_dim**-0.5 + + self.mlp_hidden_dim = int(hidden_size * mlp_ratio) + self.hidden_size = hidden_size + + # Pre-attention normalization for image tokens + self.img_pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + + # PhotonAttention module with built-in projections and norms + self.attention = PhotonAttention( + query_dim=hidden_size, + heads=num_heads, + dim_head=self.head_dim, + bias=False, + out_bias=False, + eps=1e-6, + processor=PhotonAttnProcessor2_0(), + ) + + # mlp + self.post_attention_layernorm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.gate_proj = nn.Linear(hidden_size, self.mlp_hidden_dim, bias=False) + self.up_proj = nn.Linear(hidden_size, self.mlp_hidden_dim, bias=False) + self.down_proj = nn.Linear(self.mlp_hidden_dim, hidden_size, bias=False) + self.mlp_act = nn.GELU(approximate="tanh") + + self.modulation = Modulation(hidden_size) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + image_rotary_emb: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + **kwargs: Dict[str, Any], + ) -> torch.Tensor: + r""" + Runs modulation-gated cross-attention and MLP, with residual connections. + + Args: + hidden_states (`torch.Tensor`): + Image tokens of shape `(B, L_img, hidden_size)`. + encoder_hidden_states (`torch.Tensor`): + Text tokens of shape `(B, L_txt, hidden_size)`. + temb (`torch.Tensor`): + Conditioning vector used by `Modulation` to produce scale/shift/gates, shape `(B, hidden_size)` (or + broadcastable). + image_rotary_emb (`torch.Tensor`): + Rotary positional embeddings applied inside attention. + attention_mask (`torch.Tensor`, *optional*): + Boolean mask for text tokens of shape `(B, L_txt)`, where `0` marks padding. + **kwargs: + Additional keyword arguments for API compatibility. + + Returns: + `torch.Tensor`: + Updated image tokens of shape `(B, L_img, hidden_size)`. + """ + + mod_attn, mod_mlp = self.modulation(temb) + attn_shift, attn_scale, attn_gate = mod_attn + mlp_shift, mlp_scale, mlp_gate = mod_mlp + + hidden_states_mod = (1 + attn_scale) * self.img_pre_norm(hidden_states) + attn_shift + + attn_out = self.attention( + hidden_states=hidden_states_mod, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + image_rotary_emb=image_rotary_emb, + ) + + hidden_states = hidden_states + attn_gate * attn_out + + x = (1 + mlp_scale) * self.post_attention_layernorm(hidden_states) + mlp_shift + hidden_states = hidden_states + mlp_gate * (self.down_proj(self.mlp_act(self.gate_proj(x)) * self.up_proj(x))) + return hidden_states + + +class FinalLayer(nn.Module): + r""" + Final projection layer with adaptive LayerNorm modulation. + + This layer applies a normalized and modulated transformation to input tokens and projects them into patch-level + outputs. + + Args: + hidden_size (`int`): + Dimensionality of the input tokens. + patch_size (`int`): + Size of the square image patches. + out_channels (`int`): + Number of output channels per pixel (e.g. RGB = 3). + + Forward Inputs: + x (`torch.Tensor`): + Input tokens of shape `(B, L, hidden_size)`, where `L` is the number of patches. + vec (`torch.Tensor`): + Conditioning vector of shape `(B, hidden_size)` used to generate shift and scale parameters for adaptive + LayerNorm. + + Returns: + `torch.Tensor`: + Projected patch outputs of shape `(B, L, patch_size * patch_size * out_channels)`. + """ + + def __init__(self, hidden_size: int, patch_size: int, out_channels: int): + super().__init__() + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True) + self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True)) + + def forward(self, x: torch.Tensor, vec: torch.Tensor) -> torch.Tensor: + shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1) + x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :] + x = self.linear(x) + return x + + +def img2seq(img: torch.Tensor, patch_size: int) -> torch.Tensor: + r""" + Flattens an image tensor into a sequence of non-overlapping patches. + + Args: + img (`torch.Tensor`): + Input image tensor of shape `(B, C, H, W)`. + patch_size (`int`): + Size of each square patch. Must evenly divide both `H` and `W`. + + Returns: + `torch.Tensor`: + Flattened patch sequence of shape `(B, L, C * patch_size * patch_size)`, where `L = (H // patch_size) * (W + // patch_size)` is the number of patches. + """ + return unfold(img, kernel_size=patch_size, stride=patch_size).transpose(1, 2) + + +def seq2img(seq: torch.Tensor, patch_size: int, shape: torch.Tensor) -> torch.Tensor: + r""" + Reconstructs an image tensor from a sequence of patches (inverse of `img2seq`). + + Args: + seq (`torch.Tensor`): + Patch sequence of shape `(B, L, C * patch_size * patch_size)`, where `L = (H // patch_size) * (W // + patch_size)`. + patch_size (`int`): + Size of each square patch. + shape (`tuple` or `torch.Tensor`): + The original image spatial shape `(H, W)`. If a tensor is provided, the first two values are interpreted as + height and width. + + Returns: + `torch.Tensor`: + Reconstructed image tensor of shape `(B, C, H, W)`. + """ + if isinstance(shape, tuple): + shape = shape[-2:] + elif isinstance(shape, torch.Tensor): + shape = (int(shape[0]), int(shape[1])) + else: + raise NotImplementedError(f"shape type {type(shape)} not supported") + return fold(seq.transpose(1, 2), shape, kernel_size=patch_size, stride=patch_size) + + +class PhotonTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin): + r""" + Transformer-based 2D model for text to image generation. + + Args: + in_channels (`int`, *optional*, defaults to 16): + Number of input channels in the latent image. + patch_size (`int`, *optional*, defaults to 2): + Size of the square patches used to flatten the input image. + context_in_dim (`int`, *optional*, defaults to 2304): + Dimensionality of the text conditioning input. + hidden_size (`int`, *optional*, defaults to 1792): + Dimension of the hidden representation. + mlp_ratio (`float`, *optional*, defaults to 3.5): + Expansion ratio for the hidden dimension inside MLP blocks. + num_heads (`int`, *optional*, defaults to 28): + Number of attention heads. + depth (`int`, *optional*, defaults to 16): + Number of transformer blocks. + axes_dim (`list[int]`, *optional*): + List of dimensions for each positional embedding axis. Defaults to `[32, 32]`. + theta (`int`, *optional*, defaults to 10000): + Frequency scaling factor for rotary embeddings. + time_factor (`float`, *optional*, defaults to 1000.0): + Scaling factor applied in timestep embeddings. + time_max_period (`int`, *optional*, defaults to 10000): + Maximum frequency period for timestep embeddings. + + Attributes: + pe_embedder (`EmbedND`): + Multi-axis rotary embedding generator for positional encodings. + img_in (`nn.Linear`): + Projection layer for image patch tokens. + time_in (`MLPEmbedder`): + Embedding layer for timestep embeddings. + txt_in (`nn.Linear`): + Projection layer for text conditioning. + blocks (`nn.ModuleList`): + Stack of transformer blocks (`PhotonBlock`). + final_layer (`LastLayer`): + Projection layer mapping hidden tokens back to patch outputs. + + Methods: + attn_processors: + Returns a dictionary of all attention processors in the model. + set_attn_processor(processor): + Replaces attention processors across all attention layers. + process_inputs(image_latent, txt): + Converts inputs into patch tokens, encodes text, and produces positional encodings. + compute_timestep_embedding(timestep, dtype): + Creates a timestep embedding of dimension 256, scaled and projected. + forward_transformers(image_latent, cross_attn_conditioning, timestep, time_embedding, attention_mask, + **block_kwargs): + Runs the sequence of transformer blocks over image and text tokens. + forward(image_latent, timestep, cross_attn_conditioning, micro_conditioning, cross_attn_mask=None, + attention_kwargs=None, return_dict=True): + Full forward pass from latent input to reconstructed output image. + + Returns: + `Transformer2DModelOutput` if `return_dict=True` (default), otherwise a tuple containing: + - `sample` (`torch.Tensor`): Reconstructed image of shape `(B, C, H, W)`. + """ + + config_name = "config.json" + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + in_channels: int = 16, + patch_size: int = 2, + context_in_dim: int = 2304, + hidden_size: int = 1792, + mlp_ratio: float = 3.5, + num_heads: int = 28, + depth: int = 16, + axes_dim: list = None, + theta: int = 10000, + time_factor: float = 1000.0, + time_max_period: int = 10000, + ): + super().__init__() + + if axes_dim is None: + axes_dim = [32, 32] + + # Store parameters directly + self.in_channels = in_channels + self.patch_size = patch_size + self.out_channels = self.in_channels * self.patch_size**2 + + self.time_factor = time_factor + self.time_max_period = time_max_period + + if hidden_size % num_heads != 0: + raise ValueError(f"Hidden size {hidden_size} must be divisible by num_heads {num_heads}") + + pe_dim = hidden_size // num_heads + + if sum(axes_dim) != pe_dim: + raise ValueError(f"Got {axes_dim} but expected positional dim {pe_dim}") + + self.hidden_size = hidden_size + self.num_heads = num_heads + self.pe_embedder = PhotonEmbedND(dim=pe_dim, theta=theta, axes_dim=axes_dim) + self.img_in = nn.Linear(self.in_channels * self.patch_size**2, self.hidden_size, bias=True) + self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) + self.txt_in = nn.Linear(context_in_dim, self.hidden_size) + + self.blocks = nn.ModuleList( + [ + PhotonBlock( + self.hidden_size, + self.num_heads, + mlp_ratio=mlp_ratio, + ) + for i in range(depth) + ] + ) + + self.final_layer = FinalLayer(self.hidden_size, 1, self.out_channels) + + self.gradient_checkpointing = False + + def _compute_timestep_embedding(self, timestep: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + return self.time_in( + get_timestep_embedding( + timesteps=timestep, + embedding_dim=256, + max_period=self.time_max_period, + scale=self.time_factor, + flip_sin_to_cos=True, # Match original cos, sin order + ).to(dtype) + ) + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + attention_kwargs: Optional[Dict[str, Any]] = None, + return_dict: bool = True, + ) -> Union[Tuple[torch.Tensor, ...], Transformer2DModelOutput]: + r""" + Forward pass of the PhotonTransformer2DModel. + + The latent image is split into patch tokens, combined with text conditioning, and processed through a stack of + transformer blocks modulated by the timestep. The output is reconstructed into the latent image space. + + Args: + hidden_states (`torch.Tensor`): + Input latent image tensor of shape `(B, C, H, W)`. + timestep (`torch.Tensor`): + Timestep tensor of shape `(B,)` or `(1,)`, used for temporal conditioning. + encoder_hidden_states (`torch.Tensor`): + Text conditioning tensor of shape `(B, L_txt, context_in_dim)`. + attention_mask (`torch.Tensor`, *optional*): + Boolean mask of shape `(B, L_txt)`, where `0` marks padding in the text sequence. + attention_kwargs (`dict`, *optional*): + Additional arguments passed to attention layers. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a `Transformer2DModelOutput` or a tuple. + + Returns: + `Transformer2DModelOutput` if `return_dict=True`, otherwise a tuple: + + - `sample` (`torch.Tensor`): Output latent image of shape `(B, C, H, W)`. + """ + # Process text conditioning + txt = self.txt_in(encoder_hidden_states) + + # Convert image to sequence and embed + img = img2seq(hidden_states, self.patch_size) + img = self.img_in(img) + + # Generate positional embeddings + bs, _, h, w = hidden_states.shape + img_ids = get_image_ids(bs, h, w, patch_size=self.patch_size, device=hidden_states.device) + pe = self.pe_embedder(img_ids) + + # Compute time embedding + vec = self._compute_timestep_embedding(timestep, dtype=img.dtype) + + # Apply transformer blocks + for block in self.blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + img = self._gradient_checkpointing_func( + block.__call__, + img, + txt, + vec, + pe, + attention_mask, + ) + else: + img = block( + hidden_states=img, + encoder_hidden_states=txt, + temb=vec, + image_rotary_emb=pe, + attention_mask=attention_mask, + ) + + # Final layer and convert back to image + img = self.final_layer(img, vec) + output = seq2img(img, self.patch_size, hidden_states.shape) + + if not return_dict: + return (output,) + return Transformer2DModelOutput(sample=output) diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index c438caed57..a44c92a834 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -144,6 +144,7 @@ else: "FluxKontextPipeline", "FluxKontextInpaintPipeline", ] + _import_structure["photon"] = ["PhotonPipeline"] _import_structure["audioldm"] = ["AudioLDMPipeline"] _import_structure["audioldm2"] = [ "AudioLDM2Pipeline", @@ -717,6 +718,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: StableDiffusionXLPAGPipeline, ) from .paint_by_example import PaintByExamplePipeline + from .photon import PhotonPipeline from .pia import PIAPipeline from .pixart_alpha import PixArtAlphaPipeline, PixArtSigmaPipeline from .qwenimage import ( diff --git a/src/diffusers/pipelines/photon/__init__.py b/src/diffusers/pipelines/photon/__init__.py new file mode 100644 index 0000000000..e21e31d422 --- /dev/null +++ b/src/diffusers/pipelines/photon/__init__.py @@ -0,0 +1,63 @@ +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_additional_imports = {} +_import_structure = {"pipeline_output": ["PhotonPipelineOutput"]} + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects # noqa F403 + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["pipeline_photon"] = ["PhotonPipeline"] + +# Import T5GemmaEncoder for pipeline loading compatibility +try: + if is_transformers_available(): + import transformers + from transformers.models.t5gemma.modeling_t5gemma import T5GemmaEncoder + + _additional_imports["T5GemmaEncoder"] = T5GemmaEncoder + # Patch transformers module directly for serialization + if not hasattr(transformers, "T5GemmaEncoder"): + transformers.T5GemmaEncoder = T5GemmaEncoder +except ImportError: + pass + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 + else: + from .pipeline_output import PhotonPipelineOutput + from .pipeline_photon import PhotonPipeline + +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) + + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) + for name, value in _additional_imports.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/photon/pipeline_output.py b/src/diffusers/pipelines/photon/pipeline_output.py new file mode 100644 index 0000000000..d4b0ff4629 --- /dev/null +++ b/src/diffusers/pipelines/photon/pipeline_output.py @@ -0,0 +1,35 @@ +# Copyright 2025 The Photoroom and the HuggingFace Teams. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +from typing import List, Union + +import numpy as np +import PIL.Image + +from ...utils import BaseOutput + + +@dataclass +class PhotonPipelineOutput(BaseOutput): + """ + Output class for Photon pipelines. + + Args: + images (`List[PIL.Image.Image]` or `np.ndarray`) + List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, + num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. + """ + + images: Union[List[PIL.Image.Image], np.ndarray] diff --git a/src/diffusers/pipelines/photon/pipeline_photon.py b/src/diffusers/pipelines/photon/pipeline_photon.py new file mode 100644 index 0000000000..4a10899ede --- /dev/null +++ b/src/diffusers/pipelines/photon/pipeline_photon.py @@ -0,0 +1,768 @@ +# Copyright 2025 The Photoroom and The HuggingFace Teams. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import html +import inspect +import re +import urllib.parse as ul +from typing import Callable, Dict, List, Optional, Union + +import ftfy +import torch +from transformers import ( + AutoTokenizer, + GemmaTokenizerFast, + T5TokenizerFast, +) +from transformers.models.t5gemma.modeling_t5gemma import T5GemmaEncoder + +from diffusers.image_processor import PixArtImageProcessor +from diffusers.loaders import FromSingleFileMixin, LoraLoaderMixin, TextualInversionLoaderMixin +from diffusers.models import AutoencoderDC, AutoencoderKL +from diffusers.models.transformers.transformer_photon import PhotonTransformer2DModel +from diffusers.pipelines.photon.pipeline_output import PhotonPipelineOutput +from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from diffusers.schedulers import FlowMatchEulerDiscreteScheduler +from diffusers.utils import ( + logging, + replace_example_docstring, +) +from diffusers.utils.torch_utils import randn_tensor + + +DEFAULT_RESOLUTION = 512 + +ASPECT_RATIO_256_BIN = { + "0.46": [160, 352], + "0.6": [192, 320], + "0.78": [224, 288], + "1.0": [256, 256], + "1.29": [288, 224], + "1.67": [320, 192], + "2.2": [352, 160], +} + +ASPECT_RATIO_512_BIN = { + "0.5": [352, 704], + "0.57": [384, 672], + "0.6": [384, 640], + "0.68": [416, 608], + "0.78": [448, 576], + "0.88": [480, 544], + "1.0": [512, 512], + "1.13": [544, 480], + "1.29": [576, 448], + "1.46": [608, 416], + "1.67": [640, 384], + "1.75": [672, 384], + "2.0": [704, 352], +} + +logger = logging.get_logger(__name__) + + +class TextPreprocessor: + """Text preprocessing utility for PhotonPipeline.""" + + def __init__(self): + """Initialize text preprocessor.""" + self.bad_punct_regex = re.compile( + r"[" + + "#®•©™&@·º½¾¿¡§~" + + r"\)" + + r"\(" + + r"\]" + + r"\[" + + r"\}" + + r"\{" + + r"\|" + + r"\\" + + r"\/" + + r"\*" + + r"]{1,}" + ) + + def clean_text(self, text: str) -> str: + """Clean text using comprehensive text processing logic.""" + # See Deepfloyd https://github.com/deep-floyd/IF/blob/develop/deepfloyd_if/modules/t5.py + text = str(text) + text = ul.unquote_plus(text) + text = text.strip().lower() + text = re.sub("", "person", text) + + # Remove all urls: + text = re.sub( + r"\b((?:https?|www):(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@))", + "", + text, + ) # regex for urls + + # @ + text = re.sub(r"@[\w\d]+\b", "", text) + + # 31C0—31EF CJK Strokes through 4E00—9FFF CJK Unified Ideographs + text = re.sub(r"[\u31c0-\u31ef]+", "", text) + text = re.sub(r"[\u31f0-\u31ff]+", "", text) + text = re.sub(r"[\u3200-\u32ff]+", "", text) + text = re.sub(r"[\u3300-\u33ff]+", "", text) + text = re.sub(r"[\u3400-\u4dbf]+", "", text) + text = re.sub(r"[\u4dc0-\u4dff]+", "", text) + text = re.sub(r"[\u4e00-\u9fff]+", "", text) + + # все виды тире / all types of dash --> "-" + text = re.sub( + r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+", + "-", + text, + ) + + # кавычки к одному стандарту + text = re.sub(r"[`´«»" "¨]", '"', text) + text = re.sub(r"['']", "'", text) + + # " and & + text = re.sub(r""?", "", text) + text = re.sub(r"&", "", text) + + # ip addresses: + text = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", text) + + # article ids: + text = re.sub(r"\d:\d\d\s+$", "", text) + + # \n + text = re.sub(r"\\n", " ", text) + + # "#123", "#12345..", "123456.." + text = re.sub(r"#\d{1,3}\b", "", text) + text = re.sub(r"#\d{5,}\b", "", text) + text = re.sub(r"\b\d{6,}\b", "", text) + + # filenames: + text = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", text) + + # Clean punctuation + text = re.sub(r"[\"\']{2,}", r'"', text) # """AUSVERKAUFT""" + text = re.sub(r"[\.]{2,}", r" ", text) + + text = re.sub(self.bad_punct_regex, r" ", text) # ***AUSVERKAUFT***, #AUSVERKAUFT + text = re.sub(r"\s+\.\s+", r" ", text) # " . " + + # this-is-my-cute-cat / this_is_my_cute_cat + regex2 = re.compile(r"(?:\-|\_)") + if len(re.findall(regex2, text)) > 3: + text = re.sub(regex2, " ", text) + + # Basic cleaning + text = ftfy.fix_text(text) + text = html.unescape(html.unescape(text)) + text = text.strip() + + # Clean alphanumeric patterns + text = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", text) # jc6640 + text = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", text) # jc6640vc + text = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", text) # 6640vc231 + + # Common spam patterns + text = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", text) + text = re.sub(r"(free\s)?download(\sfree)?", "", text) + text = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", text) + text = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", text) + text = re.sub(r"\bpage\s+\d+\b", "", text) + + text = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", text) # j2d1a2a... + text = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", text) + + # Final cleanup + text = re.sub(r"\b\s+\:\s+", r": ", text) + text = re.sub(r"(\D[,\./])\b", r"\1 ", text) + text = re.sub(r"\s+", " ", text) + + text.strip() + + text = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", text) + text = re.sub(r"^[\'\_,\-\:;]", r"", text) + text = re.sub(r"[\'\_,\-\:\-\+]$", r"", text) + text = re.sub(r"^\.\S+$", "", text) + + return text.strip() + + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import PhotonPipeline + + >>> # Load pipeline with from_pretrained + >>> pipe = PhotonPipeline.from_pretrained("Photoroom/photon-512-t2i-sft") + >>> pipe.to("cuda") + + >>> prompt = "A digital painting of a rusty, vintage tram on a sandy beach" + >>> image = pipe(prompt, num_inference_steps=28, guidance_scale=5.0).images[0] + >>> image.save("photon_output.png") + ``` +""" + + +class PhotonPipeline( + DiffusionPipeline, + LoraLoaderMixin, + FromSingleFileMixin, + TextualInversionLoaderMixin, +): + r""" + Pipeline for text-to-image generation using Photon Transformer. + + This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the + library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.) + + Args: + transformer ([`PhotonTransformer2DModel`]): + The Photon transformer model to denoise the encoded image latents. + scheduler ([`FlowMatchEulerDiscreteScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded image latents. + text_encoder ([`T5GemmaEncoder`]): + Text encoder model for encoding prompts. + tokenizer ([`T5TokenizerFast` or `GemmaTokenizerFast`]): + Tokenizer for the text encoder. + vae ([`AutoencoderKL`] or [`AutoencoderDC`]): + Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. + Supports both AutoencoderKL (8x compression) and AutoencoderDC (32x compression). + """ + + model_cpu_offload_seq = "text_encoder->transformer->vae" + _callback_tensor_inputs = ["latents", "prompt_embeds"] + _optional_components = ["vae"] + + def __init__( + self, + transformer: PhotonTransformer2DModel, + scheduler: FlowMatchEulerDiscreteScheduler, + text_encoder: T5GemmaEncoder, + tokenizer: Union[T5TokenizerFast, GemmaTokenizerFast, AutoTokenizer], + vae: Optional[Union[AutoencoderKL, AutoencoderDC]] = None, + default_sample_size: Optional[int] = DEFAULT_RESOLUTION, + ): + super().__init__() + + if PhotonTransformer2DModel is None: + raise ImportError( + "PhotonTransformer2DModel is not available. Please ensure the transformer_photon module is properly installed." + ) + + self.text_preprocessor = TextPreprocessor() + self.default_sample_size = default_sample_size + self._guidance_scale = 1.0 + + self.register_modules( + transformer=transformer, + scheduler=scheduler, + text_encoder=text_encoder, + tokenizer=tokenizer, + vae=vae, + ) + + self.register_to_config(default_sample_size=self.default_sample_size) + + if vae is not None: + self.image_processor = PixArtImageProcessor(vae_scale_factor=self.vae_scale_factor) + else: + self.image_processor = None + + @property + def vae_scale_factor(self): + if self.vae is None: + return 8 + if hasattr(self.vae, "spatial_compression_ratio"): + return self.vae.spatial_compression_ratio + else: # Flux VAE + return 2 ** (len(self.vae.config.block_out_channels) - 1) + + @property + def do_classifier_free_guidance(self): + """Check if classifier-free guidance is enabled based on guidance scale.""" + return self._guidance_scale > 1.0 + + @property + def guidance_scale(self): + return self._guidance_scale + + def get_default_resolution(self): + """Determine the default resolution based on the loaded VAE and config. + + Returns: + int: The default sample size (height/width) to use for generation. + """ + default_from_config = getattr(self.config, "default_sample_size", None) + if default_from_config is not None: + return default_from_config + + return DEFAULT_RESOLUTION + + def prepare_latents( + self, + batch_size: int, + num_channels_latents: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + generator: Optional[torch.Generator] = None, + latents: Optional[torch.Tensor] = None, + ): + """Prepare initial latents for the diffusion process.""" + if latents is None: + spatial_compression = self.vae_scale_factor + latent_height, latent_width = ( + height // spatial_compression, + width // spatial_compression, + ) + shape = (batch_size, num_channels_latents, latent_height, latent_width) + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + else: + latents = latents.to(device) + return latents + + def encode_prompt( + self, + prompt: Union[str, List[str]], + device: Optional[torch.device] = None, + do_classifier_free_guidance: bool = True, + negative_prompt: str = "", + num_images_per_prompt: int = 1, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + prompt_attention_mask: Optional[torch.BoolTensor] = None, + negative_prompt_attention_mask: Optional[torch.BoolTensor] = None, + ): + """Encode text prompt using standard text encoder and tokenizer, or use precomputed embeddings.""" + if device is None: + device = self._execution_device + + if prompt_embeds is None: + if isinstance(prompt, str): + prompt = [prompt] + # Encode the prompts + prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask = ( + self._encode_prompt_standard(prompt, device, do_classifier_free_guidance, negative_prompt) + ) + + # Duplicate embeddings for each generation per prompt + if num_images_per_prompt > 1: + # Repeat prompt embeddings + bs_embed, seq_len, _ = prompt_embeds.shape + 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 prompt_attention_mask is not None: + prompt_attention_mask = prompt_attention_mask.view(bs_embed, -1) + prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1) + + # Repeat negative embeddings if using CFG + if do_classifier_free_guidance and negative_prompt_embeds is not None: + bs_embed, seq_len, _ = negative_prompt_embeds.shape + negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1) + negative_prompt_embeds = negative_prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1) + + if negative_prompt_attention_mask is not None: + negative_prompt_attention_mask = negative_prompt_attention_mask.view(bs_embed, -1) + negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 1) + + return ( + prompt_embeds, + prompt_attention_mask, + negative_prompt_embeds if do_classifier_free_guidance else None, + negative_prompt_attention_mask if do_classifier_free_guidance else None, + ) + + def _tokenize_prompts(self, prompts: List[str], device: torch.device): + """Tokenize and clean prompts.""" + cleaned = [self.text_preprocessor.clean_text(text) for text in prompts] + tokens = self.tokenizer( + cleaned, + padding="max_length", + max_length=self.tokenizer.model_max_length, + truncation=True, + return_attention_mask=True, + return_tensors="pt", + ) + return tokens["input_ids"].to(device), tokens["attention_mask"].bool().to(device) + + def _encode_prompt_standard( + self, + prompt: List[str], + device: torch.device, + do_classifier_free_guidance: bool = True, + negative_prompt: str = "", + ): + """Encode prompt using standard text encoder and tokenizer with batch processing.""" + batch_size = len(prompt) + + if do_classifier_free_guidance: + if isinstance(negative_prompt, str): + negative_prompt = [negative_prompt] * batch_size + + prompts_to_encode = negative_prompt + prompt + else: + prompts_to_encode = prompt + + input_ids, attention_mask = self._tokenize_prompts(prompts_to_encode, device) + + with torch.no_grad(): + embeddings = self.text_encoder( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + )["last_hidden_state"] + + if do_classifier_free_guidance: + uncond_text_embeddings, text_embeddings = embeddings.split(batch_size, dim=0) + uncond_cross_attn_mask, cross_attn_mask = attention_mask.split(batch_size, dim=0) + else: + text_embeddings = embeddings + cross_attn_mask = attention_mask + uncond_text_embeddings = None + uncond_cross_attn_mask = None + + return text_embeddings, cross_attn_mask, uncond_text_embeddings, uncond_cross_attn_mask + + def check_inputs( + self, + prompt: Union[str, List[str]], + height: int, + width: int, + guidance_scale: float, + callback_on_step_end_tensor_inputs: Optional[List[str]] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + ): + """Check that all inputs are in correct format.""" + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + + if prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + + if prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if prompt_embeds is not None and guidance_scale > 1.0 and negative_prompt_embeds is None: + raise ValueError( + "When `prompt_embeds` is provided and `guidance_scale > 1.0`, " + "`negative_prompt_embeds` must also be provided for classifier-free guidance." + ) + + spatial_compression = self.vae_scale_factor + if height % spatial_compression != 0 or width % spatial_compression != 0: + raise ValueError( + f"`height` and `width` have to be divisible by {spatial_compression} but are {height} and {width}." + ) + + if guidance_scale < 1.0: + raise ValueError(f"guidance_scale has to be >= 1.0 but is {guidance_scale}") + + if callback_on_step_end_tensor_inputs is not None and not isinstance(callback_on_step_end_tensor_inputs, list): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be a list but is {callback_on_step_end_tensor_inputs}" + ) + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: Union[str, List[str]] = None, + negative_prompt: str = "", + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 28, + timesteps: List[int] = None, + guidance_scale: float = 4.0, + num_images_per_prompt: Optional[int] = 1, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + prompt_attention_mask: Optional[torch.BoolTensor] = None, + negative_prompt_attention_mask: Optional[torch.BoolTensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + use_resolution_binning: bool = True, + callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + ): + """ + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds` + instead. + negative_prompt (`str`, *optional*, defaults to `""`): + 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`). + height (`int`, *optional*, defaults to self.transformer.config.sample_size * self.vae_scale_factor): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to self.transformer.config.sample_size * self.vae_scale_factor): + The width in pixels of the generated image. + num_inference_steps (`int`, *optional*, defaults to 28): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + timesteps (`List[int]`, *optional*): + 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. + guidance_scale (`float`, *optional*, defaults to 4.0): + Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). + `guidance_scale` is defined as `w` of equation 2. of [Imagen + Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > + 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, + usually at the expense of lower image quality. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will be generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided and `guidance_scale > 1`, negative embeddings will be generated from an + empty string. + prompt_attention_mask (`torch.BoolTensor`, *optional*): + Pre-generated attention mask for `prompt_embeds`. If not provided, attention mask will be generated + from `prompt` input argument. + negative_prompt_attention_mask (`torch.BoolTensor`, *optional*): + Pre-generated attention mask for `negative_prompt_embeds`. If not provided and `guidance_scale > 1`, + attention mask will be generated from an empty string. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.photon.PhotonPipelineOutput`] instead of a plain tuple. + use_resolution_binning (`bool`, *optional*, defaults to `True`): + If set to `True`, the requested height and width are first mapped to the closest resolutions using + predefined aspect ratio bins. After the produced latents are decoded into images, they are resized back + to the requested resolution. Useful for generating non-square images at optimal resolutions. + callback_on_step_end (`Callable`, *optional*): + A function that calls at the end of each denoising steps during the inference. The function is called + with the following arguments: `callback_on_step_end(self, step, timestep, callback_kwargs)`. + `callback_kwargs` will include a list of all tensors as specified by + `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include tensors that are listed + in the `._callback_tensor_inputs` attribute. + + Examples: + + Returns: + [`~pipelines.photon.PhotonPipelineOutput`] or `tuple`: [`~pipelines.photon.PhotonPipelineOutput`] if + `return_dict` is True, otherwise a `tuple. When returning a tuple, the first element is a list with the + generated images. + """ + + # 0. Set height and width + default_resolution = self.get_default_resolution() + height = height or default_resolution + width = width or default_resolution + + if use_resolution_binning: + if self.image_processor is None: + raise ValueError( + "Resolution binning requires a VAE with image_processor, but VAE is not available. " + "Set use_resolution_binning=False or provide a VAE." + ) + if self.default_sample_size <= 256: + aspect_ratio_bin = ASPECT_RATIO_256_BIN + else: + aspect_ratio_bin = ASPECT_RATIO_512_BIN + + # Store original dimensions + orig_height, orig_width = height, width + # Map to closest resolution in the bin + height, width = self.image_processor.classify_height_width_bin(height, width, ratios=aspect_ratio_bin) + + # 1. Check inputs + self.check_inputs( + prompt, + height, + width, + guidance_scale, + callback_on_step_end_tensor_inputs, + prompt_embeds, + negative_prompt_embeds, + ) + + if self.vae is None and output_type not in ["latent", "pt"]: + raise ValueError( + f"VAE is required for output_type='{output_type}' but it is not available. " + "Either provide a VAE or set output_type='latent' or 'pt' to get latent outputs." + ) + + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + # Use execution device (handles offloading scenarios including group offloading) + device = self._execution_device + + self._guidance_scale = guidance_scale + + # 2. Encode input prompt + text_embeddings, cross_attn_mask, uncond_text_embeddings, uncond_cross_attn_mask = self.encode_prompt( + prompt, + device, + do_classifier_free_guidance=self.do_classifier_free_guidance, + negative_prompt=negative_prompt, + num_images_per_prompt=num_images_per_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + ) + # Expose standard names for callbacks parity + prompt_embeds = text_embeddings + negative_prompt_embeds = uncond_text_embeddings + + # 3. Prepare timesteps + if timesteps is not None: + self.scheduler.set_timesteps(timesteps=timesteps, device=device) + timesteps = self.scheduler.timesteps + num_inference_steps = len(timesteps) + else: + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = self.scheduler.timesteps + + self.num_timesteps = len(timesteps) + + # 4. Prepare latent variables + if self.vae is not None: + num_channels_latents = self.vae.config.latent_channels + else: + # When vae is None, get latent channels from transformer + num_channels_latents = self.transformer.config.in_channels + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + text_embeddings.dtype, + device, + generator, + latents, + ) + + # 5. Prepare extra step kwargs + extra_step_kwargs = {} + accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) + if accepts_eta: + extra_step_kwargs["eta"] = 0.0 + + # 6. Prepare cross-attention embeddings and masks + if self.do_classifier_free_guidance: + ca_embed = torch.cat([uncond_text_embeddings, text_embeddings], dim=0) + ca_mask = None + if cross_attn_mask is not None and uncond_cross_attn_mask is not None: + ca_mask = torch.cat([uncond_cross_attn_mask, cross_attn_mask], dim=0) + else: + ca_embed = text_embeddings + ca_mask = cross_attn_mask + + # 7. Denoising loop + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + # Duplicate latents if using classifier-free guidance + if self.do_classifier_free_guidance: + latents_in = torch.cat([latents, latents], dim=0) + # Normalize timestep for the transformer + t_cont = (t.float() / self.scheduler.config.num_train_timesteps).view(1).repeat(2).to(device) + else: + latents_in = latents + # Normalize timestep for the transformer + t_cont = (t.float() / self.scheduler.config.num_train_timesteps).view(1).to(device) + + # Forward through transformer + noise_pred = self.transformer( + hidden_states=latents_in, + timestep=t_cont, + encoder_hidden_states=ca_embed, + attention_mask=ca_mask, + return_dict=False, + )[0] + + # Apply CFG + if self.do_classifier_free_guidance: + noise_uncond, noise_text = noise_pred.chunk(2, dim=0) + noise_pred = noise_uncond + guidance_scale * (noise_text - noise_uncond) + + # Compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_on_step_end(self, i, t, callback_kwargs) + + # Call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + # 8. Post-processing + if output_type == "latent" or (output_type == "pt" and self.vae is None): + image = latents + else: + # Unscale latents for VAE (supports both AutoencoderKL and AutoencoderDC) + scaling_factor = getattr(self.vae.config, "scaling_factor", 0.18215) + shift_factor = getattr(self.vae.config, "shift_factor", 0.0) + latents = (latents / scaling_factor) + shift_factor + # Decode using VAE (AutoencoderKL or AutoencoderDC) + image = self.vae.decode(latents, return_dict=False)[0] + # Resize back to original resolution if using binning + if use_resolution_binning: + image = self.image_processor.resize_and_crop_tensor(image, orig_width, orig_height) + + # Use standard image processor for post-processing + image = self.image_processor.postprocess(image, output_type=output_type) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return PhotonPipelineOutput(images=image) diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 5d62709c28..d379a5d4a7 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -1098,6 +1098,21 @@ class ParallelConfig(metaclass=DummyObject): requires_backends(cls, ["torch"]) +class PhotonTransformer2DModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class PixArtTransformer2DModel(metaclass=DummyObject): _backends = ["torch"] diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 3244ef12ef..52c72579cd 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -1847,6 +1847,21 @@ class PaintByExamplePipeline(metaclass=DummyObject): requires_backends(cls, ["torch", "transformers"]) +class PhotonPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class PIAPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/tests/models/transformers/test_models_transformer_photon.py b/tests/models/transformers/test_models_transformer_photon.py new file mode 100644 index 0000000000..f5185245d3 --- /dev/null +++ b/tests/models/transformers/test_models_transformer_photon.py @@ -0,0 +1,83 @@ +# coding=utf-8 +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import torch + +from diffusers.models.transformers.transformer_photon import PhotonTransformer2DModel + +from ...testing_utils import enable_full_determinism, torch_device +from ..test_modeling_common import ModelTesterMixin + + +enable_full_determinism() + + +class PhotonTransformerTests(ModelTesterMixin, unittest.TestCase): + model_class = PhotonTransformer2DModel + main_input_name = "hidden_states" + uses_custom_attn_processor = True + + @property + def dummy_input(self): + return self.prepare_dummy_input() + + @property + def input_shape(self): + return (16, 16, 16) + + @property + def output_shape(self): + return (16, 16, 16) + + def prepare_dummy_input(self, height=16, width=16): + batch_size = 1 + num_latent_channels = 16 + sequence_length = 16 + embedding_dim = 1792 + + hidden_states = torch.randn((batch_size, num_latent_channels, height, width)).to(torch_device) + encoder_hidden_states = torch.randn((batch_size, sequence_length, embedding_dim)).to(torch_device) + timestep = torch.tensor([1.0]).to(torch_device).expand(batch_size) + + return { + "hidden_states": hidden_states, + "timestep": timestep, + "encoder_hidden_states": encoder_hidden_states, + } + + def prepare_init_args_and_inputs_for_common(self): + init_dict = { + "in_channels": 16, + "patch_size": 2, + "context_in_dim": 1792, + "hidden_size": 1792, + "mlp_ratio": 3.5, + "num_heads": 28, + "depth": 4, # Smaller depth for testing + "axes_dim": [32, 32], + "theta": 10_000, + } + inputs_dict = self.prepare_dummy_input() + return init_dict, inputs_dict + + def test_gradient_checkpointing_is_applied(self): + expected_set = {"PhotonTransformer2DModel"} + super().test_gradient_checkpointing_is_applied(expected_set=expected_set) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pipelines/photon/__init__.py b/tests/pipelines/photon/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/pipelines/photon/test_pipeline_photon.py b/tests/pipelines/photon/test_pipeline_photon.py new file mode 100644 index 0000000000..c29c6ce0b0 --- /dev/null +++ b/tests/pipelines/photon/test_pipeline_photon.py @@ -0,0 +1,265 @@ +import unittest + +import numpy as np +import pytest +import torch +from transformers import AutoTokenizer +from transformers.models.t5gemma.configuration_t5gemma import T5GemmaConfig, T5GemmaModuleConfig +from transformers.models.t5gemma.modeling_t5gemma import T5GemmaEncoder + +from diffusers.models import AutoencoderDC, AutoencoderKL +from diffusers.models.transformers.transformer_photon import PhotonTransformer2DModel +from diffusers.pipelines.photon.pipeline_photon import PhotonPipeline +from diffusers.schedulers import FlowMatchEulerDiscreteScheduler +from diffusers.utils import is_transformers_version + +from ..pipeline_params import TEXT_TO_IMAGE_PARAMS +from ..test_pipelines_common import PipelineTesterMixin + + +@pytest.mark.xfail( + condition=is_transformers_version(">", "4.57.1"), + reason="See https://github.com/huggingface/diffusers/pull/12456#issuecomment-3424228544", + strict=False, +) +class PhotonPipelineFastTests(PipelineTesterMixin, unittest.TestCase): + pipeline_class = PhotonPipeline + params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} + batch_params = frozenset(["prompt", "negative_prompt", "num_images_per_prompt"]) + test_xformers_attention = False + test_layerwise_casting = True + test_group_offloading = True + + @classmethod + def setUpClass(cls): + # Ensure PhotonPipeline has an _execution_device property expected by __call__ + if not isinstance(getattr(PhotonPipeline, "_execution_device", None), property): + try: + setattr(PhotonPipeline, "_execution_device", property(lambda self: torch.device("cpu"))) + except Exception: + pass + + def get_dummy_components(self): + torch.manual_seed(0) + transformer = PhotonTransformer2DModel( + patch_size=1, + in_channels=4, + context_in_dim=8, + hidden_size=8, + mlp_ratio=2.0, + num_heads=2, + depth=1, + axes_dim=[2, 2], + ) + + torch.manual_seed(0) + vae = AutoencoderKL( + sample_size=32, + in_channels=3, + out_channels=3, + block_out_channels=(4,), + layers_per_block=1, + latent_channels=4, + norm_num_groups=1, + use_quant_conv=False, + use_post_quant_conv=False, + shift_factor=0.0, + scaling_factor=1.0, + ).eval() + + torch.manual_seed(0) + scheduler = FlowMatchEulerDiscreteScheduler() + + torch.manual_seed(0) + tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/dummy-gemma") + tokenizer.model_max_length = 64 + + torch.manual_seed(0) + + encoder_params = { + "vocab_size": tokenizer.vocab_size, + "hidden_size": 8, + "intermediate_size": 16, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 4, + "max_position_embeddings": 64, + "layer_types": ["full_attention"], + "attention_bias": False, + "attention_dropout": 0.0, + "dropout_rate": 0.0, + "hidden_activation": "gelu_pytorch_tanh", + "rms_norm_eps": 1e-06, + "attn_logit_softcapping": 50.0, + "final_logit_softcapping": 30.0, + "query_pre_attn_scalar": 4, + "rope_theta": 10000.0, + "sliding_window": 4096, + } + encoder_config = T5GemmaModuleConfig(**encoder_params) + text_encoder_config = T5GemmaConfig(encoder=encoder_config, is_encoder_decoder=False, **encoder_params) + text_encoder = T5GemmaEncoder(text_encoder_config) + + return { + "transformer": transformer, + "vae": vae, + "scheduler": scheduler, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + } + + def get_dummy_inputs(self, device, seed=0): + if str(device).startswith("mps"): + generator = torch.manual_seed(seed) + else: + generator = torch.Generator(device=device).manual_seed(seed) + return { + "prompt": "", + "negative_prompt": "", + "generator": generator, + "num_inference_steps": 2, + "guidance_scale": 1.0, + "height": 32, + "width": 32, + "output_type": "pt", + "use_resolution_binning": False, + } + + def test_inference(self): + device = "cpu" + components = self.get_dummy_components() + pipe = PhotonPipeline(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + try: + pipe.register_to_config(_execution_device="cpu") + except Exception: + pass + + inputs = self.get_dummy_inputs(device) + image = pipe(**inputs)[0] + generated_image = image[0] + + self.assertEqual(generated_image.shape, (3, 32, 32)) + expected_image = torch.zeros(3, 32, 32) + max_diff = np.abs(generated_image - expected_image).max() + self.assertLessEqual(max_diff, 1e10) + + def test_callback_inputs(self): + components = self.get_dummy_components() + pipe = PhotonPipeline(**components) + pipe = pipe.to("cpu") + pipe.set_progress_bar_config(disable=None) + try: + pipe.register_to_config(_execution_device="cpu") + except Exception: + pass + self.assertTrue( + hasattr(pipe, "_callback_tensor_inputs"), + f" {PhotonPipeline} should have `_callback_tensor_inputs` that defines a list of tensor variables its callback function can use as inputs", + ) + + def callback_inputs_subset(pipe, i, t, callback_kwargs): + for tensor_name in callback_kwargs.keys(): + assert tensor_name in pipe._callback_tensor_inputs + return callback_kwargs + + def callback_inputs_all(pipe, i, t, callback_kwargs): + for tensor_name in pipe._callback_tensor_inputs: + assert tensor_name in callback_kwargs + for tensor_name in callback_kwargs.keys(): + assert tensor_name in pipe._callback_tensor_inputs + return callback_kwargs + + inputs = self.get_dummy_inputs("cpu") + + inputs["callback_on_step_end"] = callback_inputs_subset + inputs["callback_on_step_end_tensor_inputs"] = ["latents"] + _ = pipe(**inputs)[0] + + inputs["callback_on_step_end"] = callback_inputs_all + inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs + _ = pipe(**inputs)[0] + + def test_attention_slicing_forward_pass(self, expected_max_diff=1e-3): + if not self.test_attention_slicing: + return + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + for component in pipe.components.values(): + if hasattr(component, "set_default_attn_processor"): + component.set_default_attn_processor() + pipe.to("cpu") + pipe.set_progress_bar_config(disable=None) + + def to_np_local(tensor): + if isinstance(tensor, torch.Tensor): + return tensor.detach().cpu().numpy() + return tensor + + generator_device = "cpu" + inputs = self.get_dummy_inputs(generator_device) + output_without_slicing = pipe(**inputs)[0] + + pipe.enable_attention_slicing(slice_size=1) + inputs = self.get_dummy_inputs(generator_device) + output_with_slicing1 = pipe(**inputs)[0] + + pipe.enable_attention_slicing(slice_size=2) + inputs = self.get_dummy_inputs(generator_device) + output_with_slicing2 = pipe(**inputs)[0] + + max_diff1 = np.abs(to_np_local(output_with_slicing1) - to_np_local(output_without_slicing)).max() + max_diff2 = np.abs(to_np_local(output_with_slicing2) - to_np_local(output_without_slicing)).max() + self.assertLess(max(max_diff1, max_diff2), expected_max_diff) + + def test_inference_with_autoencoder_dc(self): + """Test PhotonPipeline with AutoencoderDC (DCAE) instead of AutoencoderKL.""" + device = "cpu" + + components = self.get_dummy_components() + + torch.manual_seed(0) + vae_dc = AutoencoderDC( + in_channels=3, + latent_channels=4, + attention_head_dim=2, + encoder_block_types=( + "ResBlock", + "EfficientViTBlock", + ), + decoder_block_types=( + "ResBlock", + "EfficientViTBlock", + ), + encoder_block_out_channels=(8, 8), + decoder_block_out_channels=(8, 8), + encoder_qkv_multiscales=((), (5,)), + decoder_qkv_multiscales=((), (5,)), + encoder_layers_per_block=(1, 1), + decoder_layers_per_block=(1, 1), + upsample_block_type="interpolate", + downsample_block_type="stride_conv", + decoder_norm_types="rms_norm", + decoder_act_fns="silu", + ).eval() + + components["vae"] = vae_dc + + pipe = PhotonPipeline(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + expected_scale_factor = vae_dc.spatial_compression_ratio + self.assertEqual(pipe.vae_scale_factor, expected_scale_factor) + + inputs = self.get_dummy_inputs(device) + image = pipe(**inputs)[0] + generated_image = image[0] + + self.assertEqual(generated_image.shape, (3, 32, 32)) + expected_image = torch.zeros(3, 32, 32) + max_diff = np.abs(generated_image - expected_image).max() + self.assertLessEqual(max_diff, 1e10) From 57636ad4f4567065db122bf3716cd2006a46ac77 Mon Sep 17 00:00:00 2001 From: vb Date: Tue, 21 Oct 2025 21:29:20 +0200 Subject: [PATCH 081/123] purge HF_HUB_ENABLE_HF_TRANSFER; promote Xet (#12497) * purge HF_HUB_ENABLE_HF_TRANSFER; promote Xet * purge HF_HUB_ENABLE_HF_TRANSFER; promote Xet x2 * restrict docker build test to the ones we actually use in CI. --------- Co-authored-by: YiYi Xu Co-authored-by: Sayak Paul --- .github/workflows/benchmark.yml | 2 +- .github/workflows/build_docker_images.yml | 37 +++++++++++++++---- .github/workflows/nightly_tests.yml | 2 +- .github/workflows/pr_modular_tests.yml | 2 +- .github/workflows/pr_tests.yml | 2 +- .github/workflows/pr_tests_gpu.yml | 2 +- .github/workflows/push_tests.yml | 2 +- .github/workflows/push_tests_fast.yml | 2 +- .github/workflows/push_tests_mps.yml | 2 +- docker/diffusers-doc-builder/Dockerfile | 2 +- docker/diffusers-onnxruntime-cpu/Dockerfile | 2 +- docker/diffusers-onnxruntime-cuda/Dockerfile | 5 +-- docker/diffusers-pytorch-cpu/Dockerfile | 2 +- docker/diffusers-pytorch-cuda/Dockerfile | 2 +- .../diffusers-pytorch-minimum-cuda/Dockerfile | 2 +- .../Dockerfile | 2 +- tests/quantization/torchao/README.md | 2 +- 17 files changed, 46 insertions(+), 26 deletions(-) diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index 939cce9ffb..dc3aa102be 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -7,7 +7,7 @@ on: env: DIFFUSERS_IS_CI: yes - HF_HUB_ENABLE_HF_TRANSFER: 1 + HF_XET_HIGH_PERFORMANCE: 1 HF_HOME: /mnt/cache OMP_NUM_THREADS: 8 MKL_NUM_THREADS: 8 diff --git a/.github/workflows/build_docker_images.yml b/.github/workflows/build_docker_images.yml index 1d7be0d6bc..b1af447367 100644 --- a/.github/workflows/build_docker_images.yml +++ b/.github/workflows/build_docker_images.yml @@ -42,18 +42,39 @@ jobs: CHANGED_FILES: ${{ steps.file_changes.outputs.all }} run: | echo "$CHANGED_FILES" - for FILE in $CHANGED_FILES; do + ALLOWED_IMAGES=( + diffusers-pytorch-cpu + diffusers-pytorch-cuda + diffusers-pytorch-xformers-cuda + diffusers-pytorch-minimum-cuda + diffusers-doc-builder + ) + + declare -A IMAGES_TO_BUILD=() + + for FILE in $CHANGED_FILES; do # skip anything that isn't still on disk - if [[ ! -f "$FILE" ]]; then + if [[ ! -e "$FILE" ]]; then echo "Skipping removed file $FILE" continue - fi - if [[ "$FILE" == docker/*Dockerfile ]]; then - DOCKER_PATH="${FILE%/Dockerfile}" - DOCKER_TAG=$(basename "$DOCKER_PATH") - echo "Building Docker image for $DOCKER_TAG" - docker build -t "$DOCKER_TAG" "$DOCKER_PATH" fi + + for IMAGE in "${ALLOWED_IMAGES[@]}"; do + if [[ "$FILE" == docker/${IMAGE}/* ]]; then + IMAGES_TO_BUILD["$IMAGE"]=1 + fi + done + done + + if [[ ${#IMAGES_TO_BUILD[@]} -eq 0 ]]; then + echo "No relevant Docker changes detected." + exit 0 + fi + + for IMAGE in "${!IMAGES_TO_BUILD[@]}"; do + DOCKER_PATH="docker/${IMAGE}" + echo "Building Docker image for $IMAGE" + docker build -t "$IMAGE" "$DOCKER_PATH" done if: steps.file_changes.outputs.all != '' diff --git a/.github/workflows/nightly_tests.yml b/.github/workflows/nightly_tests.yml index c8fa3a7ad9..1738efd63b 100644 --- a/.github/workflows/nightly_tests.yml +++ b/.github/workflows/nightly_tests.yml @@ -7,7 +7,7 @@ on: env: DIFFUSERS_IS_CI: yes - HF_HUB_ENABLE_HF_TRANSFER: 1 + HF_XET_HIGH_PERFORMANCE: 1 OMP_NUM_THREADS: 8 MKL_NUM_THREADS: 8 PYTEST_TIMEOUT: 600 diff --git a/.github/workflows/pr_modular_tests.yml b/.github/workflows/pr_modular_tests.yml index 83c84fbab2..7081ee518d 100644 --- a/.github/workflows/pr_modular_tests.yml +++ b/.github/workflows/pr_modular_tests.yml @@ -26,7 +26,7 @@ concurrency: env: DIFFUSERS_IS_CI: yes - HF_HUB_ENABLE_HF_TRANSFER: 1 + HF_XET_HIGH_PERFORMANCE: 1 OMP_NUM_THREADS: 4 MKL_NUM_THREADS: 4 PYTEST_TIMEOUT: 60 diff --git a/.github/workflows/pr_tests.yml b/.github/workflows/pr_tests.yml index 03205fcec6..3306ebe43e 100644 --- a/.github/workflows/pr_tests.yml +++ b/.github/workflows/pr_tests.yml @@ -22,7 +22,7 @@ concurrency: env: DIFFUSERS_IS_CI: yes - HF_HUB_ENABLE_HF_TRANSFER: 1 + HF_XET_HIGH_PERFORMANCE: 1 OMP_NUM_THREADS: 4 MKL_NUM_THREADS: 4 PYTEST_TIMEOUT: 60 diff --git a/.github/workflows/pr_tests_gpu.yml b/.github/workflows/pr_tests_gpu.yml index 900a53da94..6c208ad7ca 100644 --- a/.github/workflows/pr_tests_gpu.yml +++ b/.github/workflows/pr_tests_gpu.yml @@ -24,7 +24,7 @@ env: DIFFUSERS_IS_CI: yes OMP_NUM_THREADS: 8 MKL_NUM_THREADS: 8 - HF_HUB_ENABLE_HF_TRANSFER: 1 + HF_XET_HIGH_PERFORMANCE: 1 PYTEST_TIMEOUT: 600 PIPELINE_USAGE_CUTOFF: 1000000000 # set high cutoff so that only always-test pipelines run diff --git a/.github/workflows/push_tests.yml b/.github/workflows/push_tests.yml index a1f5e12153..a10dfcaf4e 100644 --- a/.github/workflows/push_tests.yml +++ b/.github/workflows/push_tests.yml @@ -14,7 +14,7 @@ env: DIFFUSERS_IS_CI: yes OMP_NUM_THREADS: 8 MKL_NUM_THREADS: 8 - HF_HUB_ENABLE_HF_TRANSFER: 1 + HF_XET_HIGH_PERFORMANCE: 1 PYTEST_TIMEOUT: 600 PIPELINE_USAGE_CUTOFF: 50000 diff --git a/.github/workflows/push_tests_fast.yml b/.github/workflows/push_tests_fast.yml index 6b1dd0b2d0..ae619d481c 100644 --- a/.github/workflows/push_tests_fast.yml +++ b/.github/workflows/push_tests_fast.yml @@ -18,7 +18,7 @@ env: HF_HOME: /mnt/cache OMP_NUM_THREADS: 8 MKL_NUM_THREADS: 8 - HF_HUB_ENABLE_HF_TRANSFER: 1 + HF_XET_HIGH_PERFORMANCE: 1 PYTEST_TIMEOUT: 600 RUN_SLOW: no diff --git a/.github/workflows/push_tests_mps.yml b/.github/workflows/push_tests_mps.yml index eb6c0da225..484c7a8eeb 100644 --- a/.github/workflows/push_tests_mps.yml +++ b/.github/workflows/push_tests_mps.yml @@ -8,7 +8,7 @@ env: HF_HOME: /mnt/cache OMP_NUM_THREADS: 8 MKL_NUM_THREADS: 8 - HF_HUB_ENABLE_HF_TRANSFER: 1 + HF_XET_HIGH_PERFORMANCE: 1 PYTEST_TIMEOUT: 600 RUN_SLOW: no diff --git a/docker/diffusers-doc-builder/Dockerfile b/docker/diffusers-doc-builder/Dockerfile index 8453ef4e6c..0826901e90 100644 --- a/docker/diffusers-doc-builder/Dockerfile +++ b/docker/diffusers-doc-builder/Dockerfile @@ -33,7 +33,7 @@ RUN uv pip install --no-cache-dir "git+https://github.com/huggingface/diffusers. RUN uv pip install --no-cache-dir \ accelerate \ numpy==1.26.4 \ - hf_transfer \ + hf_xet \ setuptools==69.5.1 \ bitsandbytes \ torchao \ diff --git a/docker/diffusers-onnxruntime-cpu/Dockerfile b/docker/diffusers-onnxruntime-cpu/Dockerfile index 20197937cb..7bcc6fb9f5 100644 --- a/docker/diffusers-onnxruntime-cpu/Dockerfile +++ b/docker/diffusers-onnxruntime-cpu/Dockerfile @@ -44,6 +44,6 @@ RUN python3 -m pip install --no-cache-dir --upgrade pip uv==0.1.11 && \ scipy \ tensorboard \ transformers \ - hf_transfer + hf_xet CMD ["/bin/bash"] \ No newline at end of file diff --git a/docker/diffusers-onnxruntime-cuda/Dockerfile b/docker/diffusers-onnxruntime-cuda/Dockerfile index bd1d871033..fd425d82c3 100644 --- a/docker/diffusers-onnxruntime-cuda/Dockerfile +++ b/docker/diffusers-onnxruntime-cuda/Dockerfile @@ -38,13 +38,12 @@ RUN python3.10 -m pip install --no-cache-dir --upgrade pip uv==0.1.11 && \ datasets \ hf-doc-builder \ huggingface-hub \ - hf_transfer \ + hf_xet \ Jinja2 \ librosa \ numpy==1.26.4 \ scipy \ tensorboard \ - transformers \ - hf_transfer + transformers CMD ["/bin/bash"] \ No newline at end of file diff --git a/docker/diffusers-pytorch-cpu/Dockerfile b/docker/diffusers-pytorch-cpu/Dockerfile index dc5d0fc71c..2db4ece8e2 100644 --- a/docker/diffusers-pytorch-cpu/Dockerfile +++ b/docker/diffusers-pytorch-cpu/Dockerfile @@ -31,7 +31,7 @@ RUN uv pip install --no-cache-dir "git+https://github.com/huggingface/diffusers. RUN uv pip install --no-cache-dir \ accelerate \ numpy==1.26.4 \ - hf_transfer + hf_xet RUN apt-get clean && rm -rf /var/lib/apt/lists/* && apt-get autoremove && apt-get autoclean diff --git a/docker/diffusers-pytorch-cuda/Dockerfile b/docker/diffusers-pytorch-cuda/Dockerfile index 2bdfd409b4..a700d1db72 100644 --- a/docker/diffusers-pytorch-cuda/Dockerfile +++ b/docker/diffusers-pytorch-cuda/Dockerfile @@ -44,6 +44,6 @@ RUN uv pip install --no-cache-dir \ accelerate \ numpy==1.26.4 \ pytorch-lightning \ - hf_transfer + hf_xet CMD ["/bin/bash"] diff --git a/docker/diffusers-pytorch-minimum-cuda/Dockerfile b/docker/diffusers-pytorch-minimum-cuda/Dockerfile index a2ce193f68..00d077c5da 100644 --- a/docker/diffusers-pytorch-minimum-cuda/Dockerfile +++ b/docker/diffusers-pytorch-minimum-cuda/Dockerfile @@ -47,6 +47,6 @@ RUN uv pip install --no-cache-dir \ accelerate \ numpy==1.26.4 \ pytorch-lightning \ - hf_transfer + hf_xet CMD ["/bin/bash"] diff --git a/docker/diffusers-pytorch-xformers-cuda/Dockerfile b/docker/diffusers-pytorch-xformers-cuda/Dockerfile index 1ea258bdb7..eae7eaf4fa 100644 --- a/docker/diffusers-pytorch-xformers-cuda/Dockerfile +++ b/docker/diffusers-pytorch-xformers-cuda/Dockerfile @@ -44,7 +44,7 @@ RUN uv pip install --no-cache-dir \ accelerate \ numpy==1.26.4 \ pytorch-lightning \ - hf_transfer \ + hf_xet \ xformers CMD ["/bin/bash"] diff --git a/tests/quantization/torchao/README.md b/tests/quantization/torchao/README.md index fadc529e12..373593091a 100644 --- a/tests/quantization/torchao/README.md +++ b/tests/quantization/torchao/README.md @@ -29,7 +29,7 @@ The benchmark results for Flux and CogVideoX can be found in [this](https://gith The tests, and the expected slices, were obtained from the `aws-g6e-xlarge-plus` GPU test runners. To run the slow tests, use the following command or an equivalent: ```bash -HF_HUB_ENABLE_HF_TRANSFER=1 RUN_SLOW=1 pytest -s tests/quantization/torchao/test_torchao.py::SlowTorchAoTests +HF_XET_HIGH_PERFORMANCE=1 RUN_SLOW=1 pytest -s tests/quantization/torchao/test_torchao.py::SlowTorchAoTests ``` `diffusers-cli`: From dd07b19e27b737d844f62a8107228591f8d7bca8 Mon Sep 17 00:00:00 2001 From: David Bertoin Date: Wed, 22 Oct 2025 02:09:22 +0200 Subject: [PATCH 082/123] Prx (#12525) * rename photon to prx * rename photon into prx * Revert .gitignore to state before commit b7fb0fe9d63bf766bbe3c42ac154a043796dd370 * rename photon to prx * rename photon into prx * Revert .gitignore to state before commit b7fb0fe9d63bf766bbe3c42ac154a043796dd370 * make fix-copies --- docs/source/en/_toctree.yml | 4 +- docs/source/en/api/pipelines/photon.md | 131 ------------------ docs/source/en/api/pipelines/prx.md | 131 ++++++++++++++++++ ...ffusers.py => convert_prx_to_diffusers.py} | 38 ++--- src/diffusers/__init__.py | 8 +- src/diffusers/models/__init__.py | 4 +- src/diffusers/models/transformers/__init__.py | 2 +- ...ansformer_photon.py => transformer_prx.py} | 46 +++--- src/diffusers/pipelines/__init__.py | 4 +- .../pipelines/{photon => prx}/__init__.py | 8 +- .../{photon => prx}/pipeline_output.py | 4 +- .../pipeline_prx.py} | 35 +++-- src/diffusers/utils/dummy_pt_objects.py | 30 ++-- .../dummy_torch_and_transformers_objects.py | 30 ++-- ...oton.py => test_models_transformer_prx.py} | 8 +- tests/pipelines/{photon => prx}/__init__.py | 0 .../test_pipeline_prx.py} | 26 ++-- 17 files changed, 254 insertions(+), 255 deletions(-) delete mode 100644 docs/source/en/api/pipelines/photon.md create mode 100644 docs/source/en/api/pipelines/prx.md rename scripts/{convert_photon_to_diffusers.py => convert_prx_to_diffusers.py} (91%) rename src/diffusers/models/transformers/{transformer_photon.py => transformer_prx.py} (95%) rename src/diffusers/pipelines/{photon => prx}/__init__.py (88%) rename src/diffusers/pipelines/{photon => prx}/pipeline_output.py (93%) rename src/diffusers/pipelines/{photon/pipeline_photon.py => prx/pipeline_prx.py} (96%) rename tests/models/transformers/{test_models_transformer_photon.py => test_models_transformer_prx.py} (90%) rename tests/pipelines/{photon => prx}/__init__.py (100%) rename tests/pipelines/{photon/test_pipeline_photon.py => prx/test_pipeline_prx.py} (90%) diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 3abe89437f..1e5c4fe550 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -541,12 +541,12 @@ title: PAG - local: api/pipelines/paint_by_example title: Paint by Example - - local: api/pipelines/photon - title: Photon - local: api/pipelines/pixart title: PixArt-α - local: api/pipelines/pixart_sigma title: PixArt-Σ + - local: api/pipelines/prx + title: PRX - local: api/pipelines/qwenimage title: QwenImage - local: api/pipelines/sana diff --git a/docs/source/en/api/pipelines/photon.md b/docs/source/en/api/pipelines/photon.md deleted file mode 100644 index 293e05f0fd..0000000000 --- a/docs/source/en/api/pipelines/photon.md +++ /dev/null @@ -1,131 +0,0 @@ - - -# Photon - - -Photon generates high-quality images from text using a simplified MMDIT architecture where text tokens don't update through transformer blocks. It employs flow matching with discrete scheduling for efficient sampling and uses Google's T5Gemma-2B-2B-UL2 model for multi-language text encoding. The ~1.3B parameter transformer delivers fast inference without sacrificing quality. You can choose between Flux VAE (8x compression, 16 latent channels) for balanced quality and speed or DC-AE (32x compression, 32 latent channels) for latent compression and faster processing. - -## Available models - -Photon offers multiple variants with different VAE configurations, each optimized for specific resolutions. Base models excel with detailed prompts, capturing complex compositions and subtle details. Fine-tuned models trained on the [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) improve aesthetic quality, especially with simpler prompts. - - -| Model | Resolution | Fine-tuned | Distilled | Description | Suggested prompts | Suggested parameters | Recommended dtype | -|:-----:|:-----------------:|:----------:|:----------:|:----------:|:----------:|:----------:|:----------:| -| [`Photoroom/photon-256-t2i`](https://huggingface.co/Photoroom/photon-256-t2i)| 256 | No | No | Base model pre-trained at 256 with Flux VAE|Works best with detailed prompts in natural language|28 steps, cfg=5.0| `torch.bfloat16` | -| [`Photoroom/photon-256-t2i-sft`](https://huggingface.co/Photoroom/photon-256-t2i-sft)| 512 | Yes | No | Fine-tuned on the [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) dataset with Flux VAE | Can handle less detailed prompts|28 steps, cfg=5.0| `torch.bfloat16` | -| [`Photoroom/photon-512-t2i`](https://huggingface.co/Photoroom/photon-512-t2i)| 512 | No | No | Base model pre-trained at 512 with Flux VAE |Works best with detailed prompts in natural language|28 steps, cfg=5.0| `torch.bfloat16` | -| [`Photoroom/photon-512-t2i-sft`](https://huggingface.co/Photoroom/photon-512-t2i-sft)| 512 | Yes | No | Fine-tuned on the [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) dataset with Flux VAE | Can handle less detailed prompts in natural language|28 steps, cfg=5.0| `torch.bfloat16` | -| [`Photoroom/photon-512-t2i-sft-distilled`](https://huggingface.co/Photoroom/photon-512-t2i-sft-distilled)| 512 | Yes | Yes | 8-step distilled model from [`Photoroom/photon-512-t2i-sft`](https://huggingface.co/Photoroom/photon-512-t2i-sft) | Can handle less detailed prompts in natural language|8 steps, cfg=1.0| `torch.bfloat16` | -| [`Photoroom/photon-512-t2i-dc-ae`](https://huggingface.co/Photoroom/photon-512-t2i-dc-ae)| 512 | No | No | Base model pre-trained at 512 with [Deep Compression Autoencoder (DC-AE)](https://hanlab.mit.edu/projects/dc-ae)|Works best with detailed prompts in natural language|28 steps, cfg=5.0| `torch.bfloat16` | -| [`Photoroom/photon-512-t2i-dc-ae-sft`](https://huggingface.co/Photoroom/photon-512-t2i-dc-ae-sft)| 512 | Yes | No | Fine-tuned on the [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) dataset with [Deep Compression Autoencoder (DC-AE)](https://hanlab.mit.edu/projects/dc-ae) | Can handle less detailed prompts in natural language|28 steps, cfg=5.0| `torch.bfloat16` | -| [`Photoroom/photon-512-t2i-dc-ae-sft-distilled`](https://huggingface.co/Photoroom/photon-512-t2i-dc-ae-sft-distilled)| 512 | Yes | Yes | 8-step distilled model from [`Photoroom/photon-512-t2i-dc-ae-sft-distilled`](https://huggingface.co/Photoroom/photon-512-t2i-dc-ae-sft-distilled) | Can handle less detailed prompts in natural language|8 steps, cfg=1.0| `torch.bfloat16` |s - -Refer to [this](https://huggingface.co/collections/Photoroom/photon-models-68e66254c202ebfab99ad38e) collection for more information. - -## Loading the pipeline - -Load the pipeline with [`~DiffusionPipeline.from_pretrained`]. - -```py -from diffusers.pipelines.photon import PhotonPipeline - -# Load pipeline - VAE and text encoder will be loaded from HuggingFace -pipe = PhotonPipeline.from_pretrained("Photoroom/photon-512-t2i-sft", torch_dtype=torch.bfloat16) -pipe.to("cuda") - -prompt = "A front-facing portrait of a lion the golden savanna at sunset." -image = pipe(prompt, num_inference_steps=28, guidance_scale=5.0).images[0] -image.save("photon_output.png") -``` - -### Manual Component Loading - -Load components individually to customize the pipeline for instance to use quantized models. - -```py -import torch -from diffusers.pipelines.photon import PhotonPipeline -from diffusers.models import AutoencoderKL, AutoencoderDC -from diffusers.models.transformers.transformer_photon import PhotonTransformer2DModel -from diffusers.schedulers import FlowMatchEulerDiscreteScheduler -from transformers import T5GemmaModel, GemmaTokenizerFast -from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig -from transformers import BitsAndBytesConfig as BitsAndBytesConfig - -quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True) -# Load transformer -transformer = PhotonTransformer2DModel.from_pretrained( - "checkpoints/photon-512-t2i-sft", - subfolder="transformer", - quantization_config=quant_config, - torch_dtype=torch.bfloat16, -) - -# Load scheduler -scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( - "checkpoints/photon-512-t2i-sft", subfolder="scheduler" -) - -# Load T5Gemma text encoder -t5gemma_model = T5GemmaModel.from_pretrained("google/t5gemma-2b-2b-ul2", - quantization_config=quant_config, - torch_dtype=torch.bfloat16) -text_encoder = t5gemma_model.encoder.to(dtype=torch.bfloat16) -tokenizer = GemmaTokenizerFast.from_pretrained("google/t5gemma-2b-2b-ul2") -tokenizer.model_max_length = 256 - -# Load VAE - choose either Flux VAE or DC-AE -# Flux VAE -vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", - subfolder="vae", - quantization_config=quant_config, - torch_dtype=torch.bfloat16) - -pipe = PhotonPipeline( - transformer=transformer, - scheduler=scheduler, - text_encoder=text_encoder, - tokenizer=tokenizer, - vae=vae -) -pipe.to("cuda") -``` - - -## Memory Optimization - -For memory-constrained environments: - -```py -import torch -from diffusers.pipelines.photon import PhotonPipeline - -pipe = PhotonPipeline.from_pretrained("Photoroom/photon-512-t2i-sft", torch_dtype=torch.bfloat16) -pipe.enable_model_cpu_offload() # Offload components to CPU when not in use - -# Or use sequential CPU offload for even lower memory -pipe.enable_sequential_cpu_offload() -``` - -## PhotonPipeline - -[[autodoc]] PhotonPipeline - - all - - __call__ - -## PhotonPipelineOutput - -[[autodoc]] pipelines.photon.pipeline_output.PhotonPipelineOutput diff --git a/docs/source/en/api/pipelines/prx.md b/docs/source/en/api/pipelines/prx.md new file mode 100644 index 0000000000..16670f4bfc --- /dev/null +++ b/docs/source/en/api/pipelines/prx.md @@ -0,0 +1,131 @@ + + +# PRX + + +PRX generates high-quality images from text using a simplified MMDIT architecture where text tokens don't update through transformer blocks. It employs flow matching with discrete scheduling for efficient sampling and uses Google's T5Gemma-2B-2B-UL2 model for multi-language text encoding. The ~1.3B parameter transformer delivers fast inference without sacrificing quality. You can choose between Flux VAE (8x compression, 16 latent channels) for balanced quality and speed or DC-AE (32x compression, 32 latent channels) for latent compression and faster processing. + +## Available models + +PRX offers multiple variants with different VAE configurations, each optimized for specific resolutions. Base models excel with detailed prompts, capturing complex compositions and subtle details. Fine-tuned models trained on the [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) improve aesthetic quality, especially with simpler prompts. + + +| Model | Resolution | Fine-tuned | Distilled | Description | Suggested prompts | Suggested parameters | Recommended dtype | +|:-----:|:-----------------:|:----------:|:----------:|:----------:|:----------:|:----------:|:----------:| +| [`Photoroom/prx-256-t2i`](https://huggingface.co/Photoroom/prx-256-t2i)| 256 | No | No | Base model pre-trained at 256 with Flux VAE|Works best with detailed prompts in natural language|28 steps, cfg=5.0| `torch.bfloat16` | +| [`Photoroom/prx-256-t2i-sft`](https://huggingface.co/Photoroom/prx-256-t2i-sft)| 512 | Yes | No | Fine-tuned on the [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) dataset with Flux VAE | Can handle less detailed prompts|28 steps, cfg=5.0| `torch.bfloat16` | +| [`Photoroom/prx-512-t2i`](https://huggingface.co/Photoroom/prx-512-t2i)| 512 | No | No | Base model pre-trained at 512 with Flux VAE |Works best with detailed prompts in natural language|28 steps, cfg=5.0| `torch.bfloat16` | +| [`Photoroom/prx-512-t2i-sft`](https://huggingface.co/Photoroom/prx-512-t2i-sft)| 512 | Yes | No | Fine-tuned on the [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) dataset with Flux VAE | Can handle less detailed prompts in natural language|28 steps, cfg=5.0| `torch.bfloat16` | +| [`Photoroom/prx-512-t2i-sft-distilled`](https://huggingface.co/Photoroom/prx-512-t2i-sft-distilled)| 512 | Yes | Yes | 8-step distilled model from [`Photoroom/prx-512-t2i-sft`](https://huggingface.co/Photoroom/prx-512-t2i-sft) | Can handle less detailed prompts in natural language|8 steps, cfg=1.0| `torch.bfloat16` | +| [`Photoroom/prx-512-t2i-dc-ae`](https://huggingface.co/Photoroom/prx-512-t2i-dc-ae)| 512 | No | No | Base model pre-trained at 512 with [Deep Compression Autoencoder (DC-AE)](https://hanlab.mit.edu/projects/dc-ae)|Works best with detailed prompts in natural language|28 steps, cfg=5.0| `torch.bfloat16` | +| [`Photoroom/prx-512-t2i-dc-ae-sft`](https://huggingface.co/Photoroom/prx-512-t2i-dc-ae-sft)| 512 | Yes | No | Fine-tuned on the [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) dataset with [Deep Compression Autoencoder (DC-AE)](https://hanlab.mit.edu/projects/dc-ae) | Can handle less detailed prompts in natural language|28 steps, cfg=5.0| `torch.bfloat16` | +| [`Photoroom/prx-512-t2i-dc-ae-sft-distilled`](https://huggingface.co/Photoroom/prx-512-t2i-dc-ae-sft-distilled)| 512 | Yes | Yes | 8-step distilled model from [`Photoroom/prx-512-t2i-dc-ae-sft-distilled`](https://huggingface.co/Photoroom/prx-512-t2i-dc-ae-sft-distilled) | Can handle less detailed prompts in natural language|8 steps, cfg=1.0| `torch.bfloat16` |s + +Refer to [this](https://huggingface.co/collections/Photoroom/prx-models-68e66254c202ebfab99ad38e) collection for more information. + +## Loading the pipeline + +Load the pipeline with [`~DiffusionPipeline.from_pretrained`]. + +```py +from diffusers.pipelines.prx import PRXPipeline + +# Load pipeline - VAE and text encoder will be loaded from HuggingFace +pipe = PRXPipeline.from_pretrained("Photoroom/prx-512-t2i-sft", torch_dtype=torch.bfloat16) +pipe.to("cuda") + +prompt = "A front-facing portrait of a lion the golden savanna at sunset." +image = pipe(prompt, num_inference_steps=28, guidance_scale=5.0).images[0] +image.save("prx_output.png") +``` + +### Manual Component Loading + +Load components individually to customize the pipeline for instance to use quantized models. + +```py +import torch +from diffusers.pipelines.prx import PRXPipeline +from diffusers.models import AutoencoderKL, AutoencoderDC +from diffusers.models.transformers.transformer_prx import PRXTransformer2DModel +from diffusers.schedulers import FlowMatchEulerDiscreteScheduler +from transformers import T5GemmaModel, GemmaTokenizerFast +from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig +from transformers import BitsAndBytesConfig as BitsAndBytesConfig + +quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True) +# Load transformer +transformer = PRXTransformer2DModel.from_pretrained( + "checkpoints/prx-512-t2i-sft", + subfolder="transformer", + quantization_config=quant_config, + torch_dtype=torch.bfloat16, +) + +# Load scheduler +scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + "checkpoints/prx-512-t2i-sft", subfolder="scheduler" +) + +# Load T5Gemma text encoder +t5gemma_model = T5GemmaModel.from_pretrained("google/t5gemma-2b-2b-ul2", + quantization_config=quant_config, + torch_dtype=torch.bfloat16) +text_encoder = t5gemma_model.encoder.to(dtype=torch.bfloat16) +tokenizer = GemmaTokenizerFast.from_pretrained("google/t5gemma-2b-2b-ul2") +tokenizer.model_max_length = 256 + +# Load VAE - choose either Flux VAE or DC-AE +# Flux VAE +vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", + subfolder="vae", + quantization_config=quant_config, + torch_dtype=torch.bfloat16) + +pipe = PRXPipeline( + transformer=transformer, + scheduler=scheduler, + text_encoder=text_encoder, + tokenizer=tokenizer, + vae=vae +) +pipe.to("cuda") +``` + + +## Memory Optimization + +For memory-constrained environments: + +```py +import torch +from diffusers.pipelines.prx import PRXPipeline + +pipe = PRXPipeline.from_pretrained("Photoroom/prx-512-t2i-sft", torch_dtype=torch.bfloat16) +pipe.enable_model_cpu_offload() # Offload components to CPU when not in use + +# Or use sequential CPU offload for even lower memory +pipe.enable_sequential_cpu_offload() +``` + +## PRXPipeline + +[[autodoc]] PRXPipeline + - all + - __call__ + +## PRXPipelineOutput + +[[autodoc]] pipelines.prx.pipeline_output.PRXPipelineOutput diff --git a/scripts/convert_photon_to_diffusers.py b/scripts/convert_prx_to_diffusers.py similarity index 91% rename from scripts/convert_photon_to_diffusers.py rename to scripts/convert_prx_to_diffusers.py index c66bc31418..d9bde2f34d 100644 --- a/scripts/convert_photon_to_diffusers.py +++ b/scripts/convert_prx_to_diffusers.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 """ -Script to convert Photon checkpoint from original codebase to diffusers format. +Script to convert PRX checkpoint from original codebase to diffusers format. """ import argparse @@ -13,15 +13,15 @@ from typing import Dict, Tuple import torch from safetensors.torch import save_file -from diffusers.models.transformers.transformer_photon import PhotonTransformer2DModel -from diffusers.pipelines.photon import PhotonPipeline +from diffusers.models.transformers.transformer_prx import PRXTransformer2DModel +from diffusers.pipelines.prx import PRXPipeline DEFAULT_RESOLUTION = 512 @dataclass(frozen=True) -class PhotonBase: +class PRXBase: context_in_dim: int = 2304 hidden_size: int = 1792 mlp_ratio: float = 3.5 @@ -34,22 +34,22 @@ class PhotonBase: @dataclass(frozen=True) -class PhotonFlux(PhotonBase): +class PRXFlux(PRXBase): in_channels: int = 16 patch_size: int = 2 @dataclass(frozen=True) -class PhotonDCAE(PhotonBase): +class PRXDCAE(PRXBase): in_channels: int = 32 patch_size: int = 1 def build_config(vae_type: str) -> Tuple[dict, int]: if vae_type == "flux": - cfg = PhotonFlux() + cfg = PRXFlux() elif vae_type == "dc-ae": - cfg = PhotonDCAE() + cfg = PRXDCAE() else: raise ValueError(f"Unsupported VAE type: {vae_type}. Use 'flux' or 'dc-ae'") @@ -64,7 +64,7 @@ def create_parameter_mapping(depth: int) -> dict: # Key mappings for structural changes mapping = {} - # Map old structure (layers in PhotonBlock) to new structure (layers in PhotonAttention) + # Map old structure (layers in PRXBlock) to new structure (layers in PRXAttention) for i in range(depth): # QKV projections moved to attention module mapping[f"blocks.{i}.img_qkv_proj.weight"] = f"blocks.{i}.attention.img_qkv_proj.weight" @@ -108,8 +108,8 @@ def convert_checkpoint_parameters(old_state_dict: Dict[str, torch.Tensor], depth return converted_state_dict -def create_transformer_from_checkpoint(checkpoint_path: str, config: dict) -> PhotonTransformer2DModel: - """Create and load PhotonTransformer2DModel from old checkpoint.""" +def create_transformer_from_checkpoint(checkpoint_path: str, config: dict) -> PRXTransformer2DModel: + """Create and load PRXTransformer2DModel from old checkpoint.""" print(f"Loading checkpoint from: {checkpoint_path}") @@ -137,8 +137,8 @@ def create_transformer_from_checkpoint(checkpoint_path: str, config: dict) -> Ph converted_state_dict = convert_checkpoint_parameters(state_dict, depth=model_depth) # Create transformer with config - print("Creating PhotonTransformer2DModel...") - transformer = PhotonTransformer2DModel(**config) + print("Creating PRXTransformer2DModel...") + transformer = PRXTransformer2DModel(**config) # Load state dict print("Loading converted parameters...") @@ -221,14 +221,14 @@ def create_model_index(vae_type: str, default_image_size: int, output_path: str) vae_class = "AutoencoderDC" model_index = { - "_class_name": "PhotonPipeline", + "_class_name": "PRXPipeline", "_diffusers_version": "0.31.0.dev0", "_name_or_path": os.path.basename(output_path), "default_sample_size": default_image_size, "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"], - "text_encoder": ["photon", "T5GemmaEncoder"], + "text_encoder": ["prx", "T5GemmaEncoder"], "tokenizer": ["transformers", "GemmaTokenizerFast"], - "transformer": ["diffusers", "PhotonTransformer2DModel"], + "transformer": ["diffusers", "PRXTransformer2DModel"], "vae": ["diffusers", vae_class], } @@ -275,7 +275,7 @@ def main(args): # Verify the pipeline can be loaded try: - pipeline = PhotonPipeline.from_pretrained(args.output_path) + pipeline = PRXPipeline.from_pretrained(args.output_path) print("Pipeline loaded successfully!") print(f"Transformer: {type(pipeline.transformer).__name__}") print(f"VAE: {type(pipeline.vae).__name__}") @@ -298,10 +298,10 @@ def main(args): if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Convert Photon checkpoint to diffusers format") + parser = argparse.ArgumentParser(description="Convert PRX checkpoint to diffusers format") parser.add_argument( - "--checkpoint_path", type=str, required=True, help="Path to the original Photon checkpoint (.pth file )" + "--checkpoint_path", type=str, required=True, help="Path to the original PRX checkpoint (.pth file )" ) parser.add_argument( diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index b7086d2e0c..d9c6ba9bbf 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -232,9 +232,9 @@ else: "MultiControlNetModel", "OmniGenTransformer2DModel", "ParallelConfig", - "PhotonTransformer2DModel", "PixArtTransformer2DModel", "PriorTransformer", + "PRXTransformer2DModel", "QwenImageControlNetModel", "QwenImageMultiControlNetModel", "QwenImageTransformer2DModel", @@ -516,11 +516,11 @@ else: "MusicLDMPipeline", "OmniGenPipeline", "PaintByExamplePipeline", - "PhotonPipeline", "PIAPipeline", "PixArtAlphaPipeline", "PixArtSigmaPAGPipeline", "PixArtSigmaPipeline", + "PRXPipeline", "QwenImageControlNetInpaintPipeline", "QwenImageControlNetPipeline", "QwenImageEditInpaintPipeline", @@ -928,9 +928,9 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: MultiControlNetModel, OmniGenTransformer2DModel, ParallelConfig, - PhotonTransformer2DModel, PixArtTransformer2DModel, PriorTransformer, + PRXTransformer2DModel, QwenImageControlNetModel, QwenImageMultiControlNetModel, QwenImageTransformer2DModel, @@ -1182,11 +1182,11 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: MusicLDMPipeline, OmniGenPipeline, PaintByExamplePipeline, - PhotonPipeline, PIAPipeline, PixArtAlphaPipeline, PixArtSigmaPAGPipeline, PixArtSigmaPipeline, + PRXPipeline, QwenImageControlNetInpaintPipeline, QwenImageControlNetPipeline, QwenImageEditInpaintPipeline, diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index 2151e602b2..532db76a09 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -96,7 +96,7 @@ if is_torch_available(): _import_structure["transformers.transformer_lumina2"] = ["Lumina2Transformer2DModel"] _import_structure["transformers.transformer_mochi"] = ["MochiTransformer3DModel"] _import_structure["transformers.transformer_omnigen"] = ["OmniGenTransformer2DModel"] - _import_structure["transformers.transformer_photon"] = ["PhotonTransformer2DModel"] + _import_structure["transformers.transformer_prx"] = ["PRXTransformer2DModel"] _import_structure["transformers.transformer_qwenimage"] = ["QwenImageTransformer2DModel"] _import_structure["transformers.transformer_sd3"] = ["SD3Transformer2DModel"] _import_structure["transformers.transformer_skyreels_v2"] = ["SkyReelsV2Transformer3DModel"] @@ -191,9 +191,9 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: LuminaNextDiT2DModel, MochiTransformer3DModel, OmniGenTransformer2DModel, - PhotonTransformer2DModel, PixArtTransformer2DModel, PriorTransformer, + PRXTransformer2DModel, QwenImageTransformer2DModel, SanaTransformer2DModel, SD3Transformer2DModel, diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index ab5311518b..d8c3d9b572 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -32,7 +32,7 @@ if is_torch_available(): from .transformer_lumina2 import Lumina2Transformer2DModel from .transformer_mochi import MochiTransformer3DModel from .transformer_omnigen import OmniGenTransformer2DModel - from .transformer_photon import PhotonTransformer2DModel + from .transformer_prx import PRXTransformer2DModel from .transformer_qwenimage import QwenImageTransformer2DModel from .transformer_sd3 import SD3Transformer2DModel from .transformer_skyreels_v2 import SkyReelsV2Transformer3DModel diff --git a/src/diffusers/models/transformers/transformer_photon.py b/src/diffusers/models/transformers/transformer_prx.py similarity index 95% rename from src/diffusers/models/transformers/transformer_photon.py rename to src/diffusers/models/transformers/transformer_prx.py index 6314020c1c..9b2664b9cb 100644 --- a/src/diffusers/models/transformers/transformer_photon.py +++ b/src/diffusers/models/transformers/transformer_prx.py @@ -80,9 +80,9 @@ def apply_rope(xq: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor: return xq_out.reshape(*xq.shape).type_as(xq) -class PhotonAttnProcessor2_0: +class PRXAttnProcessor2_0: r""" - Processor for implementing Photon-style attention with multi-source tokens and RoPE. Supports multiple attention + Processor for implementing PRX-style attention with multi-source tokens and RoPE. Supports multiple attention backends (Flash Attention, Sage Attention, etc.) via dispatch_attention_fn. """ @@ -91,11 +91,11 @@ class PhotonAttnProcessor2_0: def __init__(self): if not hasattr(torch.nn.functional, "scaled_dot_product_attention"): - raise ImportError("PhotonAttnProcessor2_0 requires PyTorch 2.0, please upgrade PyTorch to 2.0.") + raise ImportError("PRXAttnProcessor2_0 requires PyTorch 2.0, please upgrade PyTorch to 2.0.") def __call__( self, - attn: "PhotonAttention", + attn: "PRXAttention", hidden_states: torch.Tensor, encoder_hidden_states: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, @@ -103,10 +103,10 @@ class PhotonAttnProcessor2_0: **kwargs, ) -> torch.Tensor: """ - Apply Photon attention using PhotonAttention module. + Apply PRX attention using PRXAttention module. Args: - attn: PhotonAttention module containing projection layers + attn: PRXAttention module containing projection layers hidden_states: Image tokens [B, L_img, D] encoder_hidden_states: Text tokens [B, L_txt, D] attention_mask: Boolean mask for text tokens [B, L_txt] @@ -114,7 +114,7 @@ class PhotonAttnProcessor2_0: """ if encoder_hidden_states is None: - raise ValueError("PhotonAttnProcessor2_0 requires 'encoder_hidden_states' containing text tokens.") + raise ValueError("PRXAttnProcessor2_0 requires 'encoder_hidden_states' containing text tokens.") # Project image tokens to Q, K, V img_qkv = attn.img_qkv_proj(hidden_states) @@ -190,14 +190,14 @@ class PhotonAttnProcessor2_0: return attn_output -class PhotonAttention(nn.Module, AttentionModuleMixin): +class PRXAttention(nn.Module, AttentionModuleMixin): r""" - Photon-style attention module that handles multi-source tokens and RoPE. Similar to FluxAttention but adapted for - Photon's architecture. + PRX-style attention module that handles multi-source tokens and RoPE. Similar to FluxAttention but adapted for + PRX's architecture. """ - _default_processor_cls = PhotonAttnProcessor2_0 - _available_processors = [PhotonAttnProcessor2_0] + _default_processor_cls = PRXAttnProcessor2_0 + _available_processors = [PRXAttnProcessor2_0] def __init__( self, @@ -251,7 +251,7 @@ class PhotonAttention(nn.Module, AttentionModuleMixin): # inspired from https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/layers.py -class PhotonEmbedND(nn.Module): +class PRXEmbedND(nn.Module): r""" N-dimensional rotary positional embedding. @@ -347,7 +347,7 @@ class Modulation(nn.Module): return tuple(out[:3]), tuple(out[3:]) -class PhotonBlock(nn.Module): +class PRXBlock(nn.Module): r""" Multimodal transformer block with text–image cross-attention, modulation, and MLP. @@ -364,7 +364,7 @@ class PhotonBlock(nn.Module): Attributes: img_pre_norm (`nn.LayerNorm`): Pre-normalization applied to image tokens before attention. - attention (`PhotonAttention`): + attention (`PRXAttention`): Multi-head attention module with built-in QKV projections and normalizations for cross-attention between image and text tokens. post_attention_layernorm (`nn.LayerNorm`): @@ -400,15 +400,15 @@ class PhotonBlock(nn.Module): # Pre-attention normalization for image tokens self.img_pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - # PhotonAttention module with built-in projections and norms - self.attention = PhotonAttention( + # PRXAttention module with built-in projections and norms + self.attention = PRXAttention( query_dim=hidden_size, heads=num_heads, dim_head=self.head_dim, bias=False, out_bias=False, eps=1e-6, - processor=PhotonAttnProcessor2_0(), + processor=PRXAttnProcessor2_0(), ) # mlp @@ -557,7 +557,7 @@ def seq2img(seq: torch.Tensor, patch_size: int, shape: torch.Tensor) -> torch.Te return fold(seq.transpose(1, 2), shape, kernel_size=patch_size, stride=patch_size) -class PhotonTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin): +class PRXTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin): r""" Transformer-based 2D model for text to image generation. @@ -595,7 +595,7 @@ class PhotonTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin): txt_in (`nn.Linear`): Projection layer for text conditioning. blocks (`nn.ModuleList`): - Stack of transformer blocks (`PhotonBlock`). + Stack of transformer blocks (`PRXBlock`). final_layer (`LastLayer`): Projection layer mapping hidden tokens back to patch outputs. @@ -661,14 +661,14 @@ class PhotonTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin): self.hidden_size = hidden_size self.num_heads = num_heads - self.pe_embedder = PhotonEmbedND(dim=pe_dim, theta=theta, axes_dim=axes_dim) + self.pe_embedder = PRXEmbedND(dim=pe_dim, theta=theta, axes_dim=axes_dim) self.img_in = nn.Linear(self.in_channels * self.patch_size**2, self.hidden_size, bias=True) self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) self.txt_in = nn.Linear(context_in_dim, self.hidden_size) self.blocks = nn.ModuleList( [ - PhotonBlock( + PRXBlock( self.hidden_size, self.num_heads, mlp_ratio=mlp_ratio, @@ -702,7 +702,7 @@ class PhotonTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin): return_dict: bool = True, ) -> Union[Tuple[torch.Tensor, ...], Transformer2DModelOutput]: r""" - Forward pass of the PhotonTransformer2DModel. + Forward pass of the PRXTransformer2DModel. The latent image is split into patch tokens, combined with text conditioning, and processed through a stack of transformer blocks modulated by the timestep. The output is reconstructed into the latent image space. diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index a44c92a834..ff64958d46 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -144,7 +144,7 @@ else: "FluxKontextPipeline", "FluxKontextInpaintPipeline", ] - _import_structure["photon"] = ["PhotonPipeline"] + _import_structure["prx"] = ["PRXPipeline"] _import_structure["audioldm"] = ["AudioLDMPipeline"] _import_structure["audioldm2"] = [ "AudioLDM2Pipeline", @@ -718,9 +718,9 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: StableDiffusionXLPAGPipeline, ) from .paint_by_example import PaintByExamplePipeline - from .photon import PhotonPipeline from .pia import PIAPipeline from .pixart_alpha import PixArtAlphaPipeline, PixArtSigmaPipeline + from .prx import PRXPipeline from .qwenimage import ( QwenImageControlNetInpaintPipeline, QwenImageControlNetPipeline, diff --git a/src/diffusers/pipelines/photon/__init__.py b/src/diffusers/pipelines/prx/__init__.py similarity index 88% rename from src/diffusers/pipelines/photon/__init__.py rename to src/diffusers/pipelines/prx/__init__.py index e21e31d422..87aaefbd13 100644 --- a/src/diffusers/pipelines/photon/__init__.py +++ b/src/diffusers/pipelines/prx/__init__.py @@ -12,7 +12,7 @@ from ...utils import ( _dummy_objects = {} _additional_imports = {} -_import_structure = {"pipeline_output": ["PhotonPipelineOutput"]} +_import_structure = {"pipeline_output": ["PRXPipelineOutput"]} try: if not (is_transformers_available() and is_torch_available()): @@ -22,7 +22,7 @@ except OptionalDependencyNotAvailable: _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) else: - _import_structure["pipeline_photon"] = ["PhotonPipeline"] + _import_structure["pipeline_prx"] = ["PRXPipeline"] # Import T5GemmaEncoder for pipeline loading compatibility try: @@ -44,8 +44,8 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: except OptionalDependencyNotAvailable: from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 else: - from .pipeline_output import PhotonPipelineOutput - from .pipeline_photon import PhotonPipeline + from .pipeline_output import PRXPipelineOutput + from .pipeline_prx import PRXPipeline else: import sys diff --git a/src/diffusers/pipelines/photon/pipeline_output.py b/src/diffusers/pipelines/prx/pipeline_output.py similarity index 93% rename from src/diffusers/pipelines/photon/pipeline_output.py rename to src/diffusers/pipelines/prx/pipeline_output.py index d4b0ff4629..ea1bc9bf41 100644 --- a/src/diffusers/pipelines/photon/pipeline_output.py +++ b/src/diffusers/pipelines/prx/pipeline_output.py @@ -22,9 +22,9 @@ from ...utils import BaseOutput @dataclass -class PhotonPipelineOutput(BaseOutput): +class PRXPipelineOutput(BaseOutput): """ - Output class for Photon pipelines. + Output class for PRX pipelines. Args: images (`List[PIL.Image.Image]` or `np.ndarray`) diff --git a/src/diffusers/pipelines/photon/pipeline_photon.py b/src/diffusers/pipelines/prx/pipeline_prx.py similarity index 96% rename from src/diffusers/pipelines/photon/pipeline_photon.py rename to src/diffusers/pipelines/prx/pipeline_prx.py index 4a10899ede..a3bd3e6b45 100644 --- a/src/diffusers/pipelines/photon/pipeline_photon.py +++ b/src/diffusers/pipelines/prx/pipeline_prx.py @@ -30,9 +30,9 @@ from transformers.models.t5gemma.modeling_t5gemma import T5GemmaEncoder from diffusers.image_processor import PixArtImageProcessor from diffusers.loaders import FromSingleFileMixin, LoraLoaderMixin, TextualInversionLoaderMixin from diffusers.models import AutoencoderDC, AutoencoderKL -from diffusers.models.transformers.transformer_photon import PhotonTransformer2DModel -from diffusers.pipelines.photon.pipeline_output import PhotonPipelineOutput +from diffusers.models.transformers.transformer_prx import PRXTransformer2DModel from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from diffusers.pipelines.prx.pipeline_output import PRXPipelineOutput from diffusers.schedulers import FlowMatchEulerDiscreteScheduler from diffusers.utils import ( logging, @@ -73,7 +73,7 @@ logger = logging.get_logger(__name__) class TextPreprocessor: - """Text preprocessing utility for PhotonPipeline.""" + """Text preprocessing utility for PRXPipeline.""" def __init__(self): """Initialize text preprocessor.""" @@ -203,34 +203,34 @@ EXAMPLE_DOC_STRING = """ Examples: ```py >>> import torch - >>> from diffusers import PhotonPipeline + >>> from diffusers import PRXPipeline >>> # Load pipeline with from_pretrained - >>> pipe = PhotonPipeline.from_pretrained("Photoroom/photon-512-t2i-sft") + >>> pipe = PRXPipeline.from_pretrained("Photoroom/prx-512-t2i-sft") >>> pipe.to("cuda") >>> prompt = "A digital painting of a rusty, vintage tram on a sandy beach" >>> image = pipe(prompt, num_inference_steps=28, guidance_scale=5.0).images[0] - >>> image.save("photon_output.png") + >>> image.save("prx_output.png") ``` """ -class PhotonPipeline( +class PRXPipeline( DiffusionPipeline, LoraLoaderMixin, FromSingleFileMixin, TextualInversionLoaderMixin, ): r""" - Pipeline for text-to-image generation using Photon Transformer. + Pipeline for text-to-image generation using PRX Transformer. This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.) Args: - transformer ([`PhotonTransformer2DModel`]): - The Photon transformer model to denoise the encoded image latents. + transformer ([`PRXTransformer2DModel`]): + The PRX transformer model to denoise the encoded image latents. scheduler ([`FlowMatchEulerDiscreteScheduler`]): A scheduler to be used in combination with `transformer` to denoise the encoded image latents. text_encoder ([`T5GemmaEncoder`]): @@ -248,7 +248,7 @@ class PhotonPipeline( def __init__( self, - transformer: PhotonTransformer2DModel, + transformer: PRXTransformer2DModel, scheduler: FlowMatchEulerDiscreteScheduler, text_encoder: T5GemmaEncoder, tokenizer: Union[T5TokenizerFast, GemmaTokenizerFast, AutoTokenizer], @@ -257,9 +257,9 @@ class PhotonPipeline( ): super().__init__() - if PhotonTransformer2DModel is None: + if PRXTransformer2DModel is None: raise ImportError( - "PhotonTransformer2DModel is not available. Please ensure the transformer_photon module is properly installed." + "PRXTransformer2DModel is not available. Please ensure the transformer_prx module is properly installed." ) self.text_preprocessor = TextPreprocessor() @@ -567,7 +567,7 @@ class PhotonPipeline( The output format of the generate image. Choose between [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.photon.PhotonPipelineOutput`] instead of a plain tuple. + Whether or not to return a [`~pipelines.prx.PRXPipelineOutput`] instead of a plain tuple. use_resolution_binning (`bool`, *optional*, defaults to `True`): If set to `True`, the requested height and width are first mapped to the closest resolutions using predefined aspect ratio bins. After the produced latents are decoded into images, they are resized back @@ -585,9 +585,8 @@ class PhotonPipeline( Examples: Returns: - [`~pipelines.photon.PhotonPipelineOutput`] or `tuple`: [`~pipelines.photon.PhotonPipelineOutput`] if - `return_dict` is True, otherwise a `tuple. When returning a tuple, the first element is a list with the - generated images. + [`~pipelines.prx.PRXPipelineOutput`] or `tuple`: [`~pipelines.prx.PRXPipelineOutput`] if `return_dict` is + True, otherwise a `tuple. When returning a tuple, the first element is a list with the generated images. """ # 0. Set height and width @@ -765,4 +764,4 @@ class PhotonPipeline( if not return_dict: return (image,) - return PhotonPipelineOutput(images=image) + return PRXPipelineOutput(images=image) diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index d379a5d4a7..ecf2d7957a 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -1098,21 +1098,6 @@ class ParallelConfig(metaclass=DummyObject): requires_backends(cls, ["torch"]) -class PhotonTransformer2DModel(metaclass=DummyObject): - _backends = ["torch"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch"]) - - class PixArtTransformer2DModel(metaclass=DummyObject): _backends = ["torch"] @@ -1143,6 +1128,21 @@ class PriorTransformer(metaclass=DummyObject): requires_backends(cls, ["torch"]) +class PRXTransformer2DModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class QwenImageControlNetModel(metaclass=DummyObject): _backends = ["torch"] diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 52c72579cd..3a106c1b83 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -1847,21 +1847,6 @@ class PaintByExamplePipeline(metaclass=DummyObject): requires_backends(cls, ["torch", "transformers"]) -class PhotonPipeline(metaclass=DummyObject): - _backends = ["torch", "transformers"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch", "transformers"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers"]) - - class PIAPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] @@ -1922,6 +1907,21 @@ class PixArtSigmaPipeline(metaclass=DummyObject): requires_backends(cls, ["torch", "transformers"]) +class PRXPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class QwenImageControlNetInpaintPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/tests/models/transformers/test_models_transformer_photon.py b/tests/models/transformers/test_models_transformer_prx.py similarity index 90% rename from tests/models/transformers/test_models_transformer_photon.py rename to tests/models/transformers/test_models_transformer_prx.py index f5185245d3..1387625d5e 100644 --- a/tests/models/transformers/test_models_transformer_photon.py +++ b/tests/models/transformers/test_models_transformer_prx.py @@ -17,7 +17,7 @@ import unittest import torch -from diffusers.models.transformers.transformer_photon import PhotonTransformer2DModel +from diffusers.models.transformers.transformer_prx import PRXTransformer2DModel from ...testing_utils import enable_full_determinism, torch_device from ..test_modeling_common import ModelTesterMixin @@ -26,8 +26,8 @@ from ..test_modeling_common import ModelTesterMixin enable_full_determinism() -class PhotonTransformerTests(ModelTesterMixin, unittest.TestCase): - model_class = PhotonTransformer2DModel +class PRXTransformerTests(ModelTesterMixin, unittest.TestCase): + model_class = PRXTransformer2DModel main_input_name = "hidden_states" uses_custom_attn_processor = True @@ -75,7 +75,7 @@ class PhotonTransformerTests(ModelTesterMixin, unittest.TestCase): return init_dict, inputs_dict def test_gradient_checkpointing_is_applied(self): - expected_set = {"PhotonTransformer2DModel"} + expected_set = {"PRXTransformer2DModel"} super().test_gradient_checkpointing_is_applied(expected_set=expected_set) diff --git a/tests/pipelines/photon/__init__.py b/tests/pipelines/prx/__init__.py similarity index 100% rename from tests/pipelines/photon/__init__.py rename to tests/pipelines/prx/__init__.py diff --git a/tests/pipelines/photon/test_pipeline_photon.py b/tests/pipelines/prx/test_pipeline_prx.py similarity index 90% rename from tests/pipelines/photon/test_pipeline_photon.py rename to tests/pipelines/prx/test_pipeline_prx.py index c29c6ce0b0..46c6a5760e 100644 --- a/tests/pipelines/photon/test_pipeline_photon.py +++ b/tests/pipelines/prx/test_pipeline_prx.py @@ -8,8 +8,8 @@ from transformers.models.t5gemma.configuration_t5gemma import T5GemmaConfig, T5G from transformers.models.t5gemma.modeling_t5gemma import T5GemmaEncoder from diffusers.models import AutoencoderDC, AutoencoderKL -from diffusers.models.transformers.transformer_photon import PhotonTransformer2DModel -from diffusers.pipelines.photon.pipeline_photon import PhotonPipeline +from diffusers.models.transformers.transformer_prx import PRXTransformer2DModel +from diffusers.pipelines.prx.pipeline_prx import PRXPipeline from diffusers.schedulers import FlowMatchEulerDiscreteScheduler from diffusers.utils import is_transformers_version @@ -22,8 +22,8 @@ from ..test_pipelines_common import PipelineTesterMixin reason="See https://github.com/huggingface/diffusers/pull/12456#issuecomment-3424228544", strict=False, ) -class PhotonPipelineFastTests(PipelineTesterMixin, unittest.TestCase): - pipeline_class = PhotonPipeline +class PRXPipelineFastTests(PipelineTesterMixin, unittest.TestCase): + pipeline_class = PRXPipeline params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} batch_params = frozenset(["prompt", "negative_prompt", "num_images_per_prompt"]) test_xformers_attention = False @@ -32,16 +32,16 @@ class PhotonPipelineFastTests(PipelineTesterMixin, unittest.TestCase): @classmethod def setUpClass(cls): - # Ensure PhotonPipeline has an _execution_device property expected by __call__ - if not isinstance(getattr(PhotonPipeline, "_execution_device", None), property): + # Ensure PRXPipeline has an _execution_device property expected by __call__ + if not isinstance(getattr(PRXPipeline, "_execution_device", None), property): try: - setattr(PhotonPipeline, "_execution_device", property(lambda self: torch.device("cpu"))) + setattr(PRXPipeline, "_execution_device", property(lambda self: torch.device("cpu"))) except Exception: pass def get_dummy_components(self): torch.manual_seed(0) - transformer = PhotonTransformer2DModel( + transformer = PRXTransformer2DModel( patch_size=1, in_channels=4, context_in_dim=8, @@ -129,7 +129,7 @@ class PhotonPipelineFastTests(PipelineTesterMixin, unittest.TestCase): def test_inference(self): device = "cpu" components = self.get_dummy_components() - pipe = PhotonPipeline(**components) + pipe = PRXPipeline(**components) pipe.to(device) pipe.set_progress_bar_config(disable=None) try: @@ -148,7 +148,7 @@ class PhotonPipelineFastTests(PipelineTesterMixin, unittest.TestCase): def test_callback_inputs(self): components = self.get_dummy_components() - pipe = PhotonPipeline(**components) + pipe = PRXPipeline(**components) pipe = pipe.to("cpu") pipe.set_progress_bar_config(disable=None) try: @@ -157,7 +157,7 @@ class PhotonPipelineFastTests(PipelineTesterMixin, unittest.TestCase): pass self.assertTrue( hasattr(pipe, "_callback_tensor_inputs"), - f" {PhotonPipeline} should have `_callback_tensor_inputs` that defines a list of tensor variables its callback function can use as inputs", + f" {PRXPipeline} should have `_callback_tensor_inputs` that defines a list of tensor variables its callback function can use as inputs", ) def callback_inputs_subset(pipe, i, t, callback_kwargs): @@ -216,7 +216,7 @@ class PhotonPipelineFastTests(PipelineTesterMixin, unittest.TestCase): self.assertLess(max(max_diff1, max_diff2), expected_max_diff) def test_inference_with_autoencoder_dc(self): - """Test PhotonPipeline with AutoencoderDC (DCAE) instead of AutoencoderKL.""" + """Test PRXPipeline with AutoencoderDC (DCAE) instead of AutoencoderKL.""" device = "cpu" components = self.get_dummy_components() @@ -248,7 +248,7 @@ class PhotonPipelineFastTests(PipelineTesterMixin, unittest.TestCase): components["vae"] = vae_dc - pipe = PhotonPipeline(**components) + pipe = PRXPipeline(**components) pipe.to(device) pipe.set_progress_bar_config(disable=None) From a5a0ccf86a8b2468709f964704dd3667cbb7ac8f Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Wed, 22 Oct 2025 08:52:06 +0530 Subject: [PATCH 083/123] [core] `AutoencoderMixin` to abstract common methods (#12473) * up * correct wording. * up * up * up --- .../autoencoders/autoencoder_asym_kl.py | 7 +--- .../models/autoencoders/autoencoder_dc.py | 25 +------------ .../models/autoencoders/autoencoder_kl.py | 33 +---------------- .../autoencoders/autoencoder_kl_allegro.py | 32 +--------------- .../autoencoders/autoencoder_kl_cogvideox.py | 25 +------------ .../autoencoders/autoencoder_kl_cosmos.py | 25 +------------ .../autoencoder_kl_hunyuan_video.py | 25 +------------ .../models/autoencoders/autoencoder_kl_ltx.py | 25 +------------ .../autoencoders/autoencoder_kl_magvit.py | 25 +------------ .../autoencoders/autoencoder_kl_mochi.py | 25 +------------ .../autoencoders/autoencoder_kl_qwenimage.py | 25 +------------ .../autoencoder_kl_temporal_decoder.py | 4 +- .../models/autoencoders/autoencoder_kl_wan.py | 25 +------------ .../autoencoders/autoencoder_oobleck.py | 17 +-------- .../models/autoencoders/autoencoder_tiny.py | 33 +---------------- .../autoencoders/consistency_decoder_vae.py | 37 +------------------ src/diffusers/models/autoencoders/vae.py | 35 ++++++++++++++++++ src/diffusers/models/autoencoders/vq_model.py | 3 +- tests/models/autoencoders/testing_utils.py | 5 +++ 19 files changed, 74 insertions(+), 357 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_asym_kl.py b/src/diffusers/models/autoencoders/autoencoder_asym_kl.py index 54b1fc6771..fa49fcfe79 100644 --- a/src/diffusers/models/autoencoders/autoencoder_asym_kl.py +++ b/src/diffusers/models/autoencoders/autoencoder_asym_kl.py @@ -20,10 +20,10 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...utils.accelerate_utils import apply_forward_hook from ..modeling_outputs import AutoencoderKLOutput from ..modeling_utils import ModelMixin -from .vae import DecoderOutput, DiagonalGaussianDistribution, Encoder, MaskConditionDecoder +from .vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution, Encoder, MaskConditionDecoder -class AsymmetricAutoencoderKL(ModelMixin, ConfigMixin): +class AsymmetricAutoencoderKL(ModelMixin, AutoencoderMixin, ConfigMixin): r""" Designing a Better Asymmetric VQGAN for StableDiffusion https://huggingface.co/papers/2306.04632 . A VAE model with KL loss for encoding images into latents and decoding latent representations into images. @@ -107,9 +107,6 @@ class AsymmetricAutoencoderKL(ModelMixin, ConfigMixin): self.quant_conv = nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1) self.post_quant_conv = nn.Conv2d(latent_channels, latent_channels, 1) - self.use_slicing = False - self.use_tiling = False - self.register_to_config(block_out_channels=up_block_out_channels) self.register_to_config(force_upcast=False) diff --git a/src/diffusers/models/autoencoders/autoencoder_dc.py b/src/diffusers/models/autoencoders/autoencoder_dc.py index 783f22e97d..724ec3bb76 100644 --- a/src/diffusers/models/autoencoders/autoencoder_dc.py +++ b/src/diffusers/models/autoencoders/autoencoder_dc.py @@ -27,7 +27,7 @@ from ..attention_processor import SanaMultiscaleLinearAttention from ..modeling_utils import ModelMixin from ..normalization import RMSNorm, get_normalization from ..transformers.sana_transformer import GLUMBConv -from .vae import DecoderOutput, EncoderOutput +from .vae import AutoencoderMixin, DecoderOutput, EncoderOutput class ResBlock(nn.Module): @@ -378,7 +378,7 @@ class Decoder(nn.Module): return hidden_states -class AutoencoderDC(ModelMixin, ConfigMixin, FromOriginalModelMixin): +class AutoencoderDC(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin): r""" An Autoencoder model introduced in [DCAE](https://huggingface.co/papers/2410.10733) and used in [SANA](https://huggingface.co/papers/2410.10629). @@ -536,27 +536,6 @@ class AutoencoderDC(ModelMixin, ConfigMixin, FromOriginalModelMixin): self.tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio self.tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio - def disable_tiling(self) -> None: - r""" - Disable tiled AE decoding. If `enable_tiling` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_tiling = False - - def enable_slicing(self) -> None: - r""" - Enable sliced AE decoding. When this option is enabled, the AE will split the input tensor in slices to compute - decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - self.use_slicing = True - - def disable_slicing(self) -> None: - r""" - Disable sliced AE decoding. If `enable_slicing` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_slicing = False - def _encode(self, x: torch.Tensor) -> torch.Tensor: batch_size, num_channels, height, width = x.shape diff --git a/src/diffusers/models/autoencoders/autoencoder_kl.py b/src/diffusers/models/autoencoders/autoencoder_kl.py index d823c2fb8b..1a72aa3cfe 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl.py @@ -32,10 +32,10 @@ from ..attention_processor import ( ) from ..modeling_outputs import AutoencoderKLOutput from ..modeling_utils import ModelMixin -from .vae import Decoder, DecoderOutput, DiagonalGaussianDistribution, Encoder +from .vae import AutoencoderMixin, Decoder, DecoderOutput, DiagonalGaussianDistribution, Encoder -class AutoencoderKL(ModelMixin, ConfigMixin, FromOriginalModelMixin, PeftAdapterMixin): +class AutoencoderKL(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin, PeftAdapterMixin): r""" A VAE model with KL loss for encoding images into latents and decoding latent representations into images. @@ -138,35 +138,6 @@ class AutoencoderKL(ModelMixin, ConfigMixin, FromOriginalModelMixin, PeftAdapter self.tile_latent_min_size = int(sample_size / (2 ** (len(self.config.block_out_channels) - 1))) self.tile_overlap_factor = 0.25 - def enable_tiling(self, use_tiling: bool = True): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - self.use_tiling = use_tiling - - def disable_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.enable_tiling(False) - - def enable_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - self.use_slicing = True - - def disable_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_slicing = False - @property # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors def attn_processors(self) -> Dict[str, AttentionProcessor]: diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_allegro.py b/src/diffusers/models/autoencoders/autoencoder_kl_allegro.py index c24b8f42ac..6756586460 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_allegro.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_allegro.py @@ -28,6 +28,7 @@ from ..modeling_outputs import AutoencoderKLOutput from ..modeling_utils import ModelMixin from ..resnet import ResnetBlock2D from ..upsampling import Upsample2D +from .vae import AutoencoderMixin class AllegroTemporalConvLayer(nn.Module): @@ -673,7 +674,7 @@ class AllegroDecoder3D(nn.Module): return sample -class AutoencoderKLAllegro(ModelMixin, ConfigMixin): +class AutoencoderKLAllegro(ModelMixin, AutoencoderMixin, ConfigMixin): r""" A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. Used in [Allegro](https://github.com/rhymes-ai/Allegro). @@ -795,35 +796,6 @@ class AutoencoderKLAllegro(ModelMixin, ConfigMixin): sample_size - self.tile_overlap_w, ) - def enable_tiling(self) -> None: - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - self.use_tiling = True - - def disable_tiling(self) -> None: - r""" - Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_tiling = False - - def enable_slicing(self) -> None: - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - self.use_slicing = True - - def disable_slicing(self) -> None: - r""" - Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_slicing = False - def _encode(self, x: torch.Tensor) -> torch.Tensor: # TODO(aryan) # if self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height): diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py b/src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py index e0e9436e89..5096b725d0 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py @@ -29,7 +29,7 @@ from ..downsampling import CogVideoXDownsample3D from ..modeling_outputs import AutoencoderKLOutput from ..modeling_utils import ModelMixin from ..upsampling import CogVideoXUpsample3D -from .vae import DecoderOutput, DiagonalGaussianDistribution +from .vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -955,7 +955,7 @@ class CogVideoXDecoder3D(nn.Module): return hidden_states, new_conv_cache -class AutoencoderKLCogVideoX(ModelMixin, ConfigMixin, FromOriginalModelMixin): +class AutoencoderKLCogVideoX(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin): r""" A VAE model with KL loss for encoding images into latents and decoding latent representations into images. Used in [CogVideoX](https://github.com/THUDM/CogVideo). @@ -1124,27 +1124,6 @@ class AutoencoderKLCogVideoX(ModelMixin, ConfigMixin, FromOriginalModelMixin): self.tile_overlap_factor_height = tile_overlap_factor_height or self.tile_overlap_factor_height self.tile_overlap_factor_width = tile_overlap_factor_width or self.tile_overlap_factor_width - def disable_tiling(self) -> None: - r""" - Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_tiling = False - - def enable_slicing(self) -> None: - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - self.use_slicing = True - - def disable_slicing(self) -> None: - r""" - Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_slicing = False - def _encode(self, x: torch.Tensor) -> torch.Tensor: batch_size, num_channels, num_frames, height, width = x.shape diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_cosmos.py b/src/diffusers/models/autoencoders/autoencoder_kl_cosmos.py index 500e316ebc..b17522d1c4 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_cosmos.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_cosmos.py @@ -24,7 +24,7 @@ from ...utils import get_logger from ...utils.accelerate_utils import apply_forward_hook from ..modeling_outputs import AutoencoderKLOutput from ..modeling_utils import ModelMixin -from .vae import DecoderOutput, IdentityDistribution +from .vae import AutoencoderMixin, DecoderOutput, IdentityDistribution logger = get_logger(__name__) @@ -875,7 +875,7 @@ class CosmosDecoder3d(nn.Module): return hidden_states -class AutoencoderKLCosmos(ModelMixin, ConfigMixin): +class AutoencoderKLCosmos(ModelMixin, AutoencoderMixin, ConfigMixin): r""" Autoencoder used in [Cosmos](https://huggingface.co/papers/2501.03575). @@ -1031,27 +1031,6 @@ class AutoencoderKLCosmos(ModelMixin, ConfigMixin): self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width self.tile_sample_stride_num_frames = tile_sample_stride_num_frames or self.tile_sample_stride_num_frames - def disable_tiling(self) -> None: - r""" - Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_tiling = False - - def enable_slicing(self) -> None: - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - self.use_slicing = True - - def disable_slicing(self) -> None: - r""" - Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_slicing = False - def _encode(self, x: torch.Tensor) -> torch.Tensor: x = self.encoder(x) enc = self.quant_conv(x) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py index dc5e775f67..88b9bb507f 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py @@ -26,7 +26,7 @@ from ..activations import get_activation from ..attention_processor import Attention from ..modeling_outputs import AutoencoderKLOutput from ..modeling_utils import ModelMixin -from .vae import DecoderOutput, DiagonalGaussianDistribution +from .vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -624,7 +624,7 @@ class HunyuanVideoDecoder3D(nn.Module): return hidden_states -class AutoencoderKLHunyuanVideo(ModelMixin, ConfigMixin): +class AutoencoderKLHunyuanVideo(ModelMixin, AutoencoderMixin, ConfigMixin): r""" A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. Introduced in [HunyuanVideo](https://huggingface.co/papers/2412.03603). @@ -763,27 +763,6 @@ class AutoencoderKLHunyuanVideo(ModelMixin, ConfigMixin): self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width self.tile_sample_stride_num_frames = tile_sample_stride_num_frames or self.tile_sample_stride_num_frames - def disable_tiling(self) -> None: - r""" - Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_tiling = False - - def enable_slicing(self) -> None: - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - self.use_slicing = True - - def disable_slicing(self) -> None: - r""" - Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_slicing = False - def _encode(self, x: torch.Tensor) -> torch.Tensor: batch_size, num_channels, num_frames, height, width = x.shape diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_ltx.py b/src/diffusers/models/autoencoders/autoencoder_kl_ltx.py index 51c600a4e9..47f2081b7e 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_ltx.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_ltx.py @@ -26,7 +26,7 @@ from ..embeddings import PixArtAlphaCombinedTimestepSizeEmbeddings from ..modeling_outputs import AutoencoderKLOutput from ..modeling_utils import ModelMixin from ..normalization import RMSNorm -from .vae import DecoderOutput, DiagonalGaussianDistribution +from .vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution class LTXVideoCausalConv3d(nn.Module): @@ -1034,7 +1034,7 @@ class LTXVideoDecoder3d(nn.Module): return hidden_states -class AutoencoderKLLTXVideo(ModelMixin, ConfigMixin, FromOriginalModelMixin): +class AutoencoderKLLTXVideo(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin): r""" A VAE model with KL loss for encoding images into latents and decoding latent representations into images. Used in [LTX](https://huggingface.co/Lightricks/LTX-Video). @@ -1219,27 +1219,6 @@ class AutoencoderKLLTXVideo(ModelMixin, ConfigMixin, FromOriginalModelMixin): self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width self.tile_sample_stride_num_frames = tile_sample_stride_num_frames or self.tile_sample_stride_num_frames - def disable_tiling(self) -> None: - r""" - Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_tiling = False - - def enable_slicing(self) -> None: - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - self.use_slicing = True - - def disable_slicing(self) -> None: - r""" - Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_slicing = False - def _encode(self, x: torch.Tensor) -> torch.Tensor: batch_size, num_channels, num_frames, height, width = x.shape diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_magvit.py b/src/diffusers/models/autoencoders/autoencoder_kl_magvit.py index 43294a901f..97ca9d6692 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_magvit.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_magvit.py @@ -26,7 +26,7 @@ from ...utils.accelerate_utils import apply_forward_hook from ..activations import get_activation from ..modeling_outputs import AutoencoderKLOutput from ..modeling_utils import ModelMixin -from .vae import DecoderOutput, DiagonalGaussianDistribution +from .vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -663,7 +663,7 @@ class EasyAnimateDecoder(nn.Module): return hidden_states -class AutoencoderKLMagvit(ModelMixin, ConfigMixin): +class AutoencoderKLMagvit(ModelMixin, AutoencoderMixin, ConfigMixin): r""" A VAE model with KL loss for encoding images into latents and decoding latent representations into images. This model is used in [EasyAnimate](https://huggingface.co/papers/2405.18991). @@ -805,27 +805,6 @@ class AutoencoderKLMagvit(ModelMixin, ConfigMixin): self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width self.tile_sample_stride_num_frames = tile_sample_stride_num_frames or self.tile_sample_stride_num_frames - def disable_tiling(self) -> None: - r""" - Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_tiling = False - - def enable_slicing(self) -> None: - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - self.use_slicing = True - - def disable_slicing(self) -> None: - r""" - Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_slicing = False - @apply_forward_hook def _encode( self, x: torch.Tensor, return_dict: bool = True diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_mochi.py b/src/diffusers/models/autoencoders/autoencoder_kl_mochi.py index 404d2f6d86..3ded9a0a54 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_mochi.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_mochi.py @@ -27,7 +27,7 @@ from ..attention_processor import Attention, MochiVaeAttnProcessor2_0 from ..modeling_outputs import AutoencoderKLOutput from ..modeling_utils import ModelMixin from .autoencoder_kl_cogvideox import CogVideoXCausalConv3d -from .vae import DecoderOutput, DiagonalGaussianDistribution +from .vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -657,7 +657,7 @@ class MochiDecoder3D(nn.Module): return hidden_states, new_conv_cache -class AutoencoderKLMochi(ModelMixin, ConfigMixin): +class AutoencoderKLMochi(ModelMixin, AutoencoderMixin, ConfigMixin): r""" A VAE model with KL loss for encoding images into latents and decoding latent representations into images. Used in [Mochi 1 preview](https://github.com/genmoai/models). @@ -818,27 +818,6 @@ class AutoencoderKLMochi(ModelMixin, ConfigMixin): self.tile_sample_stride_height = tile_sample_stride_height or self.tile_sample_stride_height self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width - def disable_tiling(self) -> None: - r""" - Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_tiling = False - - def enable_slicing(self) -> None: - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - self.use_slicing = True - - def disable_slicing(self) -> None: - r""" - Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_slicing = False - def _enable_framewise_encoding(self): r""" Enables the framewise VAE encoding implementation with past latent padding. By default, Diffusers uses the diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py index 9872cf0968..14db6aeb61 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py @@ -31,7 +31,7 @@ from ...utils.accelerate_utils import apply_forward_hook from ..activations import get_activation from ..modeling_outputs import AutoencoderKLOutput from ..modeling_utils import ModelMixin -from .vae import DecoderOutput, DiagonalGaussianDistribution +from .vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -663,7 +663,7 @@ class QwenImageDecoder3d(nn.Module): return x -class AutoencoderKLQwenImage(ModelMixin, ConfigMixin, FromOriginalModelMixin): +class AutoencoderKLQwenImage(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin): r""" A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. @@ -763,27 +763,6 @@ class AutoencoderKLQwenImage(ModelMixin, ConfigMixin, FromOriginalModelMixin): self.tile_sample_stride_height = tile_sample_stride_height or self.tile_sample_stride_height self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width - def disable_tiling(self) -> None: - r""" - Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_tiling = False - - def enable_slicing(self) -> None: - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - self.use_slicing = True - - def disable_slicing(self) -> None: - r""" - Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_slicing = False - def clear_cache(self): def _count_conv3d(model): count = 0 diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py b/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py index cf46e52564..ab76254d19 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py @@ -23,7 +23,7 @@ from ..attention_processor import CROSS_ATTENTION_PROCESSORS, AttentionProcessor from ..modeling_outputs import AutoencoderKLOutput from ..modeling_utils import ModelMixin from ..unets.unet_3d_blocks import MidBlockTemporalDecoder, UpBlockTemporalDecoder -from .vae import DecoderOutput, DiagonalGaussianDistribution, Encoder +from .vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution, Encoder class TemporalDecoder(nn.Module): @@ -135,7 +135,7 @@ class TemporalDecoder(nn.Module): return sample -class AutoencoderKLTemporalDecoder(ModelMixin, ConfigMixin): +class AutoencoderKLTemporalDecoder(ModelMixin, AutoencoderMixin, ConfigMixin): r""" A VAE model with KL loss for encoding images into latents and decoding latent representations into images. diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_wan.py b/src/diffusers/models/autoencoders/autoencoder_kl_wan.py index f95c4cf374..a3aaa84650 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_wan.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_wan.py @@ -25,7 +25,7 @@ from ...utils.accelerate_utils import apply_forward_hook from ..activations import get_activation from ..modeling_outputs import AutoencoderKLOutput from ..modeling_utils import ModelMixin -from .vae import DecoderOutput, DiagonalGaussianDistribution +from .vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -951,7 +951,7 @@ def unpatchify(x, patch_size): return x -class AutoencoderKLWan(ModelMixin, ConfigMixin, FromOriginalModelMixin): +class AutoencoderKLWan(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin): r""" A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. Introduced in [Wan 2.1]. @@ -1110,27 +1110,6 @@ class AutoencoderKLWan(ModelMixin, ConfigMixin, FromOriginalModelMixin): self.tile_sample_stride_height = tile_sample_stride_height or self.tile_sample_stride_height self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width - def disable_tiling(self) -> None: - r""" - Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_tiling = False - - def enable_slicing(self) -> None: - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - self.use_slicing = True - - def disable_slicing(self) -> None: - r""" - Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_slicing = False - def clear_cache(self): # Use cached conv counts for decoder and encoder to avoid re-iterating modules each call self._conv_num = self._cached_conv_counts["decoder"] diff --git a/src/diffusers/models/autoencoders/autoencoder_oobleck.py b/src/diffusers/models/autoencoders/autoencoder_oobleck.py index a10b616b4e..d832645592 100644 --- a/src/diffusers/models/autoencoders/autoencoder_oobleck.py +++ b/src/diffusers/models/autoencoders/autoencoder_oobleck.py @@ -25,6 +25,7 @@ from ...utils import BaseOutput from ...utils.accelerate_utils import apply_forward_hook from ...utils.torch_utils import randn_tensor from ..modeling_utils import ModelMixin +from .vae import AutoencoderMixin class Snake1d(nn.Module): @@ -291,7 +292,7 @@ class OobleckDecoder(nn.Module): return hidden_state -class AutoencoderOobleck(ModelMixin, ConfigMixin): +class AutoencoderOobleck(ModelMixin, AutoencoderMixin, ConfigMixin): r""" An autoencoder for encoding waveforms into latents and decoding latent representations into waveforms. First introduced in Stable Audio. @@ -356,20 +357,6 @@ class AutoencoderOobleck(ModelMixin, ConfigMixin): self.use_slicing = False - def enable_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - self.use_slicing = True - - def disable_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_slicing = False - @apply_forward_hook def encode( self, x: torch.Tensor, return_dict: bool = True diff --git a/src/diffusers/models/autoencoders/autoencoder_tiny.py b/src/diffusers/models/autoencoders/autoencoder_tiny.py index 3e2b28606e..b9ac713d73 100644 --- a/src/diffusers/models/autoencoders/autoencoder_tiny.py +++ b/src/diffusers/models/autoencoders/autoencoder_tiny.py @@ -22,7 +22,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...utils import BaseOutput from ...utils.accelerate_utils import apply_forward_hook from ..modeling_utils import ModelMixin -from .vae import DecoderOutput, DecoderTiny, EncoderTiny +from .vae import AutoencoderMixin, DecoderOutput, DecoderTiny, EncoderTiny @dataclass @@ -38,7 +38,7 @@ class AutoencoderTinyOutput(BaseOutput): latents: torch.Tensor -class AutoencoderTiny(ModelMixin, ConfigMixin): +class AutoencoderTiny(ModelMixin, AutoencoderMixin, ConfigMixin): r""" A tiny distilled VAE model for encoding images into latents and decoding latent representations into images. @@ -162,35 +162,6 @@ class AutoencoderTiny(ModelMixin, ConfigMixin): """[0, 1] -> raw latents""" return x.sub(self.latent_shift).mul(2 * self.latent_magnitude) - def enable_slicing(self) -> None: - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - self.use_slicing = True - - def disable_slicing(self) -> None: - r""" - Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_slicing = False - - def enable_tiling(self, use_tiling: bool = True) -> None: - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - self.use_tiling = use_tiling - - def disable_tiling(self) -> None: - r""" - Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.enable_tiling(False) - def _tiled_encode(self, x: torch.Tensor) -> torch.Tensor: r"""Encode a batch of images using a tiled encoder. diff --git a/src/diffusers/models/autoencoders/consistency_decoder_vae.py b/src/diffusers/models/autoencoders/consistency_decoder_vae.py index b3017a8780..0a6258fed3 100644 --- a/src/diffusers/models/autoencoders/consistency_decoder_vae.py +++ b/src/diffusers/models/autoencoders/consistency_decoder_vae.py @@ -32,7 +32,7 @@ from ..attention_processor import ( ) from ..modeling_utils import ModelMixin from ..unets.unet_2d import UNet2DModel -from .vae import DecoderOutput, DiagonalGaussianDistribution, Encoder +from .vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution, Encoder @dataclass @@ -49,7 +49,7 @@ class ConsistencyDecoderVAEOutput(BaseOutput): latent_dist: "DiagonalGaussianDistribution" -class ConsistencyDecoderVAE(ModelMixin, ConfigMixin): +class ConsistencyDecoderVAE(ModelMixin, AutoencoderMixin, ConfigMixin): r""" The consistency decoder used with DALL-E 3. @@ -167,39 +167,6 @@ class ConsistencyDecoderVAE(ModelMixin, ConfigMixin): self.tile_latent_min_size = int(sample_size / (2 ** (len(self.config.block_out_channels) - 1))) self.tile_overlap_factor = 0.25 - # Copied from diffusers.models.autoencoders.autoencoder_kl.AutoencoderKL.enable_tiling - def enable_tiling(self, use_tiling: bool = True): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - self.use_tiling = use_tiling - - # Copied from diffusers.models.autoencoders.autoencoder_kl.AutoencoderKL.disable_tiling - def disable_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.enable_tiling(False) - - # Copied from diffusers.models.autoencoders.autoencoder_kl.AutoencoderKL.enable_slicing - def enable_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - self.use_slicing = True - - # Copied from diffusers.models.autoencoders.autoencoder_kl.AutoencoderKL.disable_slicing - def disable_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_slicing = False - @property # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors def attn_processors(self) -> Dict[str, AttentionProcessor]: diff --git a/src/diffusers/models/autoencoders/vae.py b/src/diffusers/models/autoencoders/vae.py index 1d74d4f472..7b17196125 100644 --- a/src/diffusers/models/autoencoders/vae.py +++ b/src/diffusers/models/autoencoders/vae.py @@ -894,3 +894,38 @@ class DecoderTiny(nn.Module): # scale image from [0, 1] to [-1, 1] to match diffusers convention return x.mul(2).sub(1) + + +class AutoencoderMixin: + def enable_tiling(self): + r""" + Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to + compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow + processing larger images. + """ + if not hasattr(self, "use_tiling"): + raise NotImplementedError(f"Tiling doesn't seem to be implemented for {self.__class__.__name__}.") + self.use_tiling = True + + def disable_tiling(self): + r""" + Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing + decoding in one step. + """ + self.use_tiling = False + + def enable_slicing(self): + r""" + Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to + compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. + """ + if not hasattr(self, "use_slicing"): + raise NotImplementedError(f"Slicing doesn't seem to be implemented for {self.__class__.__name__}.") + self.use_slicing = True + + def disable_slicing(self): + r""" + Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing + decoding in one step. + """ + self.use_slicing = False diff --git a/src/diffusers/models/autoencoders/vq_model.py b/src/diffusers/models/autoencoders/vq_model.py index c1094e62f7..82436473df 100644 --- a/src/diffusers/models/autoencoders/vq_model.py +++ b/src/diffusers/models/autoencoders/vq_model.py @@ -22,6 +22,7 @@ from ...utils import BaseOutput from ...utils.accelerate_utils import apply_forward_hook from ..autoencoders.vae import Decoder, DecoderOutput, Encoder, VectorQuantizer from ..modeling_utils import ModelMixin +from .vae import AutoencoderMixin @dataclass @@ -37,7 +38,7 @@ class VQEncoderOutput(BaseOutput): latents: torch.Tensor -class VQModel(ModelMixin, ConfigMixin): +class VQModel(ModelMixin, AutoencoderMixin, ConfigMixin): r""" A VQ-VAE model for decoding latent representations. diff --git a/tests/models/autoencoders/testing_utils.py b/tests/models/autoencoders/testing_utils.py index cf1f10a4a5..8ae362ac2e 100644 --- a/tests/models/autoencoders/testing_utils.py +++ b/tests/models/autoencoders/testing_utils.py @@ -57,6 +57,9 @@ class AutoencoderTesterMixin: torch.manual_seed(0) model = self.model_class(**init_dict).to(torch_device) + if not hasattr(model, "use_tiling"): + pytest.skip(f"Skipping test as {self.model_class.__name__} doesn't support tiling.") + inputs_dict.update({"return_dict": False}) _ = inputs_dict.pop("generator", None) accepts_generator = self._accepts_generator(model) @@ -102,6 +105,8 @@ class AutoencoderTesterMixin: torch.manual_seed(0) model = self.model_class(**init_dict).to(torch_device) + if not hasattr(model, "use_slicing"): + pytest.skip(f"Skipping test as {self.model_class.__name__} doesn't support tiling.") inputs_dict.update({"return_dict": False}) _ = inputs_dict.pop("generator", None) From a0a51eb0987e62bee877dcc54c11fcf1c7a892cc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=81lvaro=20Somoza?= Date: Wed, 22 Oct 2025 13:32:47 -0300 Subject: [PATCH 084/123] Kandinsky5 No cfg fix (#12527) fix --- .../pipelines/kandinsky5/pipeline_kandinsky.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py index a1122a8256..2b977a5a36 100644 --- a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py +++ b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py @@ -744,11 +744,13 @@ class Kandinsky5T2VPipeline(DiffusionPipeline, KandinskyLoraLoaderMixin): ) if negative_prompt_embeds_qwen is None: - negative_prompt_embeds_qwen, negative_prompt_embeds_clip, negative_cu_seqlens = self.encode_prompt( - prompt=negative_prompt, - max_sequence_length=max_sequence_length, - device=device, - dtype=dtype, + negative_prompt_embeds_qwen, negative_prompt_embeds_clip, negative_prompt_cu_seqlens = ( + self.encode_prompt( + prompt=negative_prompt, + max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, + ) ) # 4. Prepare timesteps @@ -780,8 +782,8 @@ class Kandinsky5T2VPipeline(DiffusionPipeline, KandinskyLoraLoaderMixin): text_rope_pos = torch.arange(prompt_cu_seqlens.diff().max().item(), device=device) negative_text_rope_pos = ( - torch.arange(negative_cu_seqlens.diff().max().item(), device=device) - if negative_cu_seqlens is not None + torch.arange(negative_prompt_cu_seqlens.diff().max().item(), device=device) + if negative_prompt_cu_seqlens is not None else None ) From bec2d8eaea97f6c0f033433d9e17abe06e0525c8 Mon Sep 17 00:00:00 2001 From: YiYi Xu Date: Wed, 22 Oct 2025 07:53:13 -1000 Subject: [PATCH 085/123] Fix: Add _skip_keys for AutoencoderKLWan (#12523) add --- .../models/autoencoders/autoencoder_kl_wan.py | 27 ++++++++++--------- src/diffusers/models/modeling_utils.py | 1 + .../pipelines/pipeline_loading_utils.py | 6 ++++- 3 files changed, 21 insertions(+), 13 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_wan.py b/src/diffusers/models/autoencoders/autoencoder_kl_wan.py index a3aaa84650..dae26f8086 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_wan.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_wan.py @@ -453,14 +453,14 @@ class WanMidBlock(nn.Module): def forward(self, x, feat_cache=None, feat_idx=[0]): # First residual block - x = self.resnets[0](x, feat_cache, feat_idx) + x = self.resnets[0](x, feat_cache=feat_cache, feat_idx=feat_idx) # Process through attention and residual blocks for attn, resnet in zip(self.attentions, self.resnets[1:]): if attn is not None: x = attn(x) - x = resnet(x, feat_cache, feat_idx) + x = resnet(x, feat_cache=feat_cache, feat_idx=feat_idx) return x @@ -494,9 +494,9 @@ class WanResidualDownBlock(nn.Module): def forward(self, x, feat_cache=None, feat_idx=[0]): x_copy = x.clone() for resnet in self.resnets: - x = resnet(x, feat_cache, feat_idx) + x = resnet(x, feat_cache=feat_cache, feat_idx=feat_idx) if self.downsampler is not None: - x = self.downsampler(x, feat_cache, feat_idx) + x = self.downsampler(x, feat_cache=feat_cache, feat_idx=feat_idx) return x + self.avg_shortcut(x_copy) @@ -598,12 +598,12 @@ class WanEncoder3d(nn.Module): ## downsamples for layer in self.down_blocks: if feat_cache is not None: - x = layer(x, feat_cache, feat_idx) + x = layer(x, feat_cache=feat_cache, feat_idx=feat_idx) else: x = layer(x) ## middle - x = self.mid_block(x, feat_cache, feat_idx) + x = self.mid_block(x, feat_cache=feat_cache, feat_idx=feat_idx) ## head x = self.norm_out(x) @@ -694,13 +694,13 @@ class WanResidualUpBlock(nn.Module): for resnet in self.resnets: if feat_cache is not None: - x = resnet(x, feat_cache, feat_idx) + x = resnet(x, feat_cache=feat_cache, feat_idx=feat_idx) else: x = resnet(x) if self.upsampler is not None: if feat_cache is not None: - x = self.upsampler(x, feat_cache, feat_idx) + x = self.upsampler(x, feat_cache=feat_cache, feat_idx=feat_idx) else: x = self.upsampler(x) @@ -767,13 +767,13 @@ class WanUpBlock(nn.Module): """ for resnet in self.resnets: if feat_cache is not None: - x = resnet(x, feat_cache, feat_idx) + x = resnet(x, feat_cache=feat_cache, feat_idx=feat_idx) else: x = resnet(x) if self.upsamplers is not None: if feat_cache is not None: - x = self.upsamplers[0](x, feat_cache, feat_idx) + x = self.upsamplers[0](x, feat_cache=feat_cache, feat_idx=feat_idx) else: x = self.upsamplers[0](x) return x @@ -885,11 +885,11 @@ class WanDecoder3d(nn.Module): x = self.conv_in(x) ## middle - x = self.mid_block(x, feat_cache, feat_idx) + x = self.mid_block(x, feat_cache=feat_cache, feat_idx=feat_idx) ## upsamples for up_block in self.up_blocks: - x = up_block(x, feat_cache, feat_idx, first_chunk=first_chunk) + x = up_block(x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk) ## head x = self.norm_out(x) @@ -961,6 +961,9 @@ class AutoencoderKLWan(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalMo """ _supports_gradient_checkpointing = False + # keys toignore when AlignDeviceHook moves inputs/outputs between devices + # these are shared mutable state modified in-place + _skip_keys = ["feat_cache", "feat_idx"] @register_to_config def __init__( diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index ce26785f63..91daca1ad8 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -251,6 +251,7 @@ class ModelMixin(torch.nn.Module, PushToHubMixin): _repeated_blocks = [] _parallel_config = None _cp_plan = None + _skip_keys = None def __init__(self): super().__init__() diff --git a/src/diffusers/pipelines/pipeline_loading_utils.py b/src/diffusers/pipelines/pipeline_loading_utils.py index dd542145d3..2169700cea 100644 --- a/src/diffusers/pipelines/pipeline_loading_utils.py +++ b/src/diffusers/pipelines/pipeline_loading_utils.py @@ -866,6 +866,9 @@ def load_sub_model( # remove hooks remove_hook_from_module(loaded_sub_model, recurse=True) needs_offloading_to_cpu = device_map[""] == "cpu" + skip_keys = None + if hasattr(loaded_sub_model, "_skip_keys") and loaded_sub_model._skip_keys is not None: + skip_keys = loaded_sub_model._skip_keys if needs_offloading_to_cpu: dispatch_model( @@ -874,9 +877,10 @@ def load_sub_model( device_map=device_map, force_hooks=True, main_device=0, + skip_keys=skip_keys, ) else: - dispatch_model(loaded_sub_model, device_map=device_map, force_hooks=True) + dispatch_model(loaded_sub_model, device_map=device_map, force_hooks=True, skip_keys=skip_keys) return loaded_sub_model From 4c723d8ec318b5fd266afaf14ba37afdefc967df Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 23 Oct 2025 00:15:47 +0530 Subject: [PATCH 086/123] [CI] xfail the test_wuerstchen_prior test (#12530) xfail the test_wuerstchen_prior test --- .../pipelines/stable_cascade/test_stable_cascade_prior.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/pipelines/stable_cascade/test_stable_cascade_prior.py b/tests/pipelines/stable_cascade/test_stable_cascade_prior.py index f8267186db..0bc821b7e6 100644 --- a/tests/pipelines/stable_cascade/test_stable_cascade_prior.py +++ b/tests/pipelines/stable_cascade/test_stable_cascade_prior.py @@ -17,11 +17,13 @@ import gc import unittest import numpy as np +import pytest import torch from transformers import CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer from diffusers import DDPMWuerstchenScheduler, StableCascadePriorPipeline from diffusers.models import StableCascadeUNet +from diffusers.utils import is_transformers_version from diffusers.utils.import_utils import is_peft_available from ...testing_utils import ( @@ -154,6 +156,11 @@ class StableCascadePriorPipelineFastTests(PipelineTesterMixin, unittest.TestCase } return inputs + @pytest.mark.xfail( + condition=is_transformers_version(">=", "4.57.1"), + reason="Test fails with the latest transformers version", + strict=False, + ) def test_wuerstchen_prior(self): device = "cpu" From ccdd96ca5231a59617661d60d5dd822cd97281fc Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 23 Oct 2025 15:09:41 +0530 Subject: [PATCH 087/123] [tests] Test attention backends (#12388) * add a lightweight test suite for attention backends. * up * up * Apply suggestions from code review * formatting --- tests/others/test_attention_backends.py | 144 ++++++++++++++++++++++++ 1 file changed, 144 insertions(+) create mode 100644 tests/others/test_attention_backends.py diff --git a/tests/others/test_attention_backends.py b/tests/others/test_attention_backends.py new file mode 100644 index 0000000000..42cdcd56f7 --- /dev/null +++ b/tests/others/test_attention_backends.py @@ -0,0 +1,144 @@ +""" +This test suite exists for the maintainers currently. It's not run in our CI at the moment. + +Once attention backends become more mature, we can consider including this in our CI. + +To run this test suite: + +```bash +export RUN_ATTENTION_BACKEND_TESTS=yes +export DIFFUSERS_ENABLE_HUB_KERNELS=yes + +pytest tests/others/test_attention_backends.py +``` + +Tests were conducted on an H100 with PyTorch 2.8.0 (CUDA 12.9). Slices for the compilation tests in +"native" variants were obtained with a torch nightly version (2.10.0.dev20250924+cu128). +""" + +import os + +import pytest +import torch + + +pytestmark = pytest.mark.skipif( + os.getenv("RUN_ATTENTION_BACKEND_TESTS", "false") == "false", reason="Feature not mature enough." +) +from diffusers import FluxPipeline # noqa: E402 +from diffusers.utils import is_torch_version # noqa: E402 + + +# fmt: off +FORWARD_CASES = [ + ("flash_hub", None), + ( + "_flash_3_hub", + torch.tensor([0.0820, 0.0859, 0.0938, 0.1016, 0.0977, 0.0996, 0.1016, 0.1016, 0.2188, 0.2246, 0.2344, 0.2480, 0.2539, 0.2480, 0.2441, 0.2715], dtype=torch.bfloat16), + ), + ( + "native", + torch.tensor([0.0820, 0.0859, 0.0938, 0.1016, 0.0957, 0.0996, 0.0996, 0.1016, 0.2188, 0.2266, 0.2363, 0.2500, 0.2539, 0.2480, 0.2461, 0.2734], dtype=torch.bfloat16) + ), + ( + "_native_cudnn", + torch.tensor([0.0781, 0.0840, 0.0879, 0.0957, 0.0898, 0.0957, 0.0957, 0.0977, 0.2168, 0.2246, 0.2324, 0.2500, 0.2539, 0.2480, 0.2441, 0.2695], dtype=torch.bfloat16), + ), +] + +COMPILE_CASES = [ + ("flash_hub", None, True), + ( + "_flash_3_hub", + torch.tensor([0.0410, 0.0410, 0.0449, 0.0508, 0.0508, 0.0605, 0.0625, 0.0605, 0.2344, 0.2461, 0.2578, 0.2734, 0.2852, 0.2812, 0.2773, 0.3047], dtype=torch.bfloat16), + True, + ), + ( + "native", + torch.tensor([0.0410, 0.0410, 0.0449, 0.0508, 0.0508, 0.0605, 0.0605, 0.0605, 0.2344, 0.2461, 0.2578, 0.2773, 0.2871, 0.2832, 0.2773, 0.3066], dtype=torch.bfloat16), + True, + ), + ( + "_native_cudnn", + torch.tensor([0.0410, 0.0410, 0.0430, 0.0508, 0.0488, 0.0586, 0.0605, 0.0586, 0.2344, 0.2461, 0.2578, 0.2773, 0.2871, 0.2832, 0.2793, 0.3086], dtype=torch.bfloat16), + True, + ), +] +# fmt: on + +INFER_KW = { + "prompt": "dance doggo dance", + "height": 256, + "width": 256, + "num_inference_steps": 2, + "guidance_scale": 3.5, + "max_sequence_length": 128, + "output_type": "pt", +} + + +def _backend_is_probably_supported(pipe, name: str): + try: + pipe.transformer.set_attention_backend(name) + return pipe, True + except Exception: + return False + + +def _check_if_slices_match(output, expected_slice): + img = output.images.detach().cpu() + generated_slice = img.flatten() + generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) + assert torch.allclose(generated_slice, expected_slice, atol=1e-4) + + +@pytest.fixture(scope="session") +def device(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for these tests.") + return torch.device("cuda:0") + + +@pytest.fixture(scope="session") +def pipe(device): + repo_id = "black-forest-labs/FLUX.1-dev" + pipe = FluxPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16).to(device) + pipe.set_progress_bar_config(disable=True) + return pipe + + +@pytest.mark.parametrize("backend_name,expected_slice", FORWARD_CASES, ids=[c[0] for c in FORWARD_CASES]) +def test_forward(pipe, backend_name, expected_slice): + out = _backend_is_probably_supported(pipe, backend_name) + if isinstance(out, bool): + pytest.xfail(f"Backend '{backend_name}' not supported in this environment.") + + modified_pipe = out[0] + out = modified_pipe(**INFER_KW, generator=torch.manual_seed(0)) + _check_if_slices_match(out, expected_slice) + + +@pytest.mark.parametrize( + "backend_name,expected_slice,error_on_recompile", + COMPILE_CASES, + ids=[c[0] for c in COMPILE_CASES], +) +def test_forward_with_compile(pipe, backend_name, expected_slice, error_on_recompile): + if "native" in backend_name and error_on_recompile and not is_torch_version(">=", "2.9.0"): + pytest.xfail(f"Test with {backend_name=} is compatible with a higher version of torch.") + + out = _backend_is_probably_supported(pipe, backend_name) + if isinstance(out, bool): + pytest.xfail(f"Backend '{backend_name}' not supported in this environment.") + + modified_pipe = out[0] + modified_pipe.transformer.compile(fullgraph=True) + + torch.compiler.reset() + with ( + torch._inductor.utils.fresh_inductor_cache(), + torch._dynamo.config.patch(error_on_recompile=error_on_recompile), + ): + out = modified_pipe(**INFER_KW, generator=torch.manual_seed(0)) + + _check_if_slices_match(out, expected_slice) From 85eb5056724f776e122627f72b524252bfc5e9be Mon Sep 17 00:00:00 2001 From: kaixuanliu Date: Thu, 23 Oct 2025 18:47:22 +0800 Subject: [PATCH 088/123] fix CI bug for kandinsky3_img2img case (#12474) * fix CI bug for kandinsky3_img2img case Signed-off-by: Liu, Kaixuan * update code Signed-off-by: Liu, Kaixuan --------- Signed-off-by: Liu, Kaixuan --- .../pipelines/kandinsky3/pipeline_kandinsky3_img2img.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py b/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py index c7b8022c22..73c2688975 100644 --- a/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py +++ b/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py @@ -113,7 +113,7 @@ class Kandinsky3Img2ImgPipeline(DiffusionPipeline, StableDiffusionLoraLoaderMixi negative_prompt=None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, - _cut_context=False, + _cut_context=True, attention_mask: Optional[torch.Tensor] = None, negative_attention_mask: Optional[torch.Tensor] = None, ): From 74b5fed43484135cd80496ab7b3ffc7936fb1ac5 Mon Sep 17 00:00:00 2001 From: Aishwarya Badlani <41635755+Aishwarya0811@users.noreply.github.com> Date: Thu, 23 Oct 2025 15:48:07 +0500 Subject: [PATCH 089/123] Fix MPS compatibility in get_1d_sincos_pos_embed_from_grid #12432 (#12449) * Fix MPS compatibility in get_1d_sincos_pos_embed_from_grid #12432 * Fix trailing whitespace in docstring * Apply style fixes --------- Co-authored-by: Dhruv Nair Co-authored-by: github-actions[bot] --- src/diffusers/models/embeddings.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/src/diffusers/models/embeddings.py b/src/diffusers/models/embeddings.py index b51f5d7aec..37fc412adc 100644 --- a/src/diffusers/models/embeddings.py +++ b/src/diffusers/models/embeddings.py @@ -319,13 +319,17 @@ def get_2d_sincos_pos_embed_from_grid(embed_dim, grid, output_type="np"): return emb -def get_1d_sincos_pos_embed_from_grid(embed_dim, pos, output_type="np", flip_sin_to_cos=False): +def get_1d_sincos_pos_embed_from_grid(embed_dim, pos, output_type="np", flip_sin_to_cos=False, dtype=None): """ This function generates 1D positional embeddings from a grid. Args: embed_dim (`int`): The embedding dimension `D` pos (`torch.Tensor`): 1D tensor of positions with shape `(M,)` + output_type (`str`, *optional*, defaults to `"np"`): Output type. Use `"pt"` for PyTorch tensors. + flip_sin_to_cos (`bool`, *optional*, defaults to `False`): Whether to flip sine and cosine embeddings. + dtype (`torch.dtype`, *optional*): Data type for frequency calculations. If `None`, defaults to + `torch.float32` on MPS devices (which don't support `torch.float64`) and `torch.float64` on other devices. Returns: `torch.Tensor`: Sinusoidal positional embeddings of shape `(M, D)`. @@ -341,7 +345,11 @@ def get_1d_sincos_pos_embed_from_grid(embed_dim, pos, output_type="np", flip_sin if embed_dim % 2 != 0: raise ValueError("embed_dim must be divisible by 2") - omega = torch.arange(embed_dim // 2, device=pos.device, dtype=torch.float64) + # Auto-detect appropriate dtype if not specified + if dtype is None: + dtype = torch.float32 if pos.device.type == "mps" else torch.float64 + + omega = torch.arange(embed_dim // 2, device=pos.device, dtype=dtype) omega /= embed_dim / 2.0 omega = 1.0 / 10000**omega # (D/2,) From 9c3b58dcf16ebd027fd3d85ec703ad5a142b1e1e Mon Sep 17 00:00:00 2001 From: Dhruv Nair Date: Thu, 23 Oct 2025 16:22:07 +0530 Subject: [PATCH 090/123] Handle deprecated transformer classes (#12517) * update * update * update --- .../pipelines/pipeline_loading_utils.py | 15 ++++++ src/diffusers/utils/__init__.py | 2 +- src/diffusers/utils/deprecation_utils.py | 48 +++++++++++++++++++ 3 files changed, 64 insertions(+), 1 deletion(-) diff --git a/src/diffusers/pipelines/pipeline_loading_utils.py b/src/diffusers/pipelines/pipeline_loading_utils.py index 2169700cea..00ef996185 100644 --- a/src/diffusers/pipelines/pipeline_loading_utils.py +++ b/src/diffusers/pipelines/pipeline_loading_utils.py @@ -33,6 +33,7 @@ from ..utils import ( ONNX_WEIGHTS_NAME, SAFETENSORS_WEIGHTS_NAME, WEIGHTS_NAME, + _maybe_remap_transformers_class, deprecate, get_class_from_dynamic_module, is_accelerate_available, @@ -356,6 +357,11 @@ def maybe_raise_or_warn( """Simple helper method to raise or warn in case incorrect module has been passed""" if not is_pipeline_module: library = importlib.import_module(library_name) + + # Handle deprecated Transformers classes + if library_name == "transformers": + class_name = _maybe_remap_transformers_class(class_name) or class_name + class_obj = getattr(library, class_name) class_candidates = {c: getattr(library, c, None) for c in importable_classes.keys()} @@ -390,6 +396,11 @@ def simple_get_class_obj(library_name, class_name): class_obj = getattr(pipeline_module, class_name) else: library = importlib.import_module(library_name) + + # Handle deprecated Transformers classes + if library_name == "transformers": + class_name = _maybe_remap_transformers_class(class_name) or class_name + class_obj = getattr(library, class_name) return class_obj @@ -416,6 +427,10 @@ def get_class_obj_and_candidates( # else we just import it from the library. library = importlib.import_module(library_name) + # Handle deprecated Transformers classes + if library_name == "transformers": + class_name = _maybe_remap_transformers_class(class_name) or class_name + class_obj = getattr(library, class_name) class_candidates = {c: getattr(library, c, None) for c in importable_classes.keys()} diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py index 63932221b2..d8e1a55401 100644 --- a/src/diffusers/utils/__init__.py +++ b/src/diffusers/utils/__init__.py @@ -38,7 +38,7 @@ from .constants import ( WEIGHTS_INDEX_NAME, WEIGHTS_NAME, ) -from .deprecation_utils import deprecate +from .deprecation_utils import _maybe_remap_transformers_class, deprecate from .doc_utils import replace_example_docstring from .dynamic_modules_utils import get_class_from_dynamic_module from .export_utils import export_to_gif, export_to_obj, export_to_ply, export_to_video diff --git a/src/diffusers/utils/deprecation_utils.py b/src/diffusers/utils/deprecation_utils.py index 4f001b3047..d76623541b 100644 --- a/src/diffusers/utils/deprecation_utils.py +++ b/src/diffusers/utils/deprecation_utils.py @@ -4,6 +4,54 @@ from typing import Any, Dict, Optional, Union from packaging import version +from ..utils import logging + + +logger = logging.get_logger(__name__) + +# Mapping for deprecated Transformers classes to their replacements +# This is used to handle models that reference deprecated class names in their configs +# Reference: https://github.com/huggingface/transformers/issues/40822 +# Format: { +# "DeprecatedClassName": { +# "new_class": "NewClassName", +# "transformers_version": (">=", "5.0.0"), # (operation, version) tuple +# } +# } +_TRANSFORMERS_CLASS_REMAPPING = { + "CLIPFeatureExtractor": { + "new_class": "CLIPImageProcessor", + "transformers_version": (">", "4.57.0"), + }, +} + + +def _maybe_remap_transformers_class(class_name: str) -> Optional[str]: + """ + Check if a Transformers class should be remapped to a newer version. + + Args: + class_name: The name of the class to check + + Returns: + The new class name if remapping should occur, None otherwise + """ + if class_name not in _TRANSFORMERS_CLASS_REMAPPING: + return None + + from .import_utils import is_transformers_version + + mapping = _TRANSFORMERS_CLASS_REMAPPING[class_name] + operation, required_version = mapping["transformers_version"] + + # Only remap if the transformers version meets the requirement + if is_transformers_version(operation, required_version): + new_class = mapping["new_class"] + logger.warning(f"{class_name} appears to have been deprecated in transformers. Using {new_class} instead.") + return mapping["new_class"] + + return None + def deprecate(*args, take_from: Optional[Union[Dict, Any]] = None, standard_warn=True, stacklevel=2): from .. import __version__ From bc4039886d4a1163b6f99912a17d4b82ad00adce Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Fri, 24 Oct 2025 12:00:02 +0530 Subject: [PATCH 091/123] fix constants.py to user `upper()` (#12479) --- src/diffusers/utils/constants.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/utils/constants.py b/src/diffusers/utils/constants.py index 8b4d76f3cb..42a53e1810 100644 --- a/src/diffusers/utils/constants.py +++ b/src/diffusers/utils/constants.py @@ -45,7 +45,7 @@ DIFFUSERS_ATTN_BACKEND = os.getenv("DIFFUSERS_ATTN_BACKEND", "native") DIFFUSERS_ATTN_CHECKS = os.getenv("DIFFUSERS_ATTN_CHECKS", "0") in ENV_VARS_TRUE_VALUES DEFAULT_HF_PARALLEL_LOADING_WORKERS = 8 HF_ENABLE_PARALLEL_LOADING = os.environ.get("HF_ENABLE_PARALLEL_LOADING", "").upper() in ENV_VARS_TRUE_VALUES -DIFFUSERS_DISABLE_REMOTE_CODE = os.getenv("DIFFUSERS_DISABLE_REMOTE_CODE", "false").lower() in ENV_VARS_TRUE_VALUES +DIFFUSERS_DISABLE_REMOTE_CODE = os.getenv("DIFFUSERS_DISABLE_REMOTE_CODE", "false").upper() in ENV_VARS_TRUE_VALUES DIFFUSERS_ENABLE_HUB_KERNELS = os.environ.get("DIFFUSERS_ENABLE_HUB_KERNELS", "").upper() in ENV_VARS_TRUE_VALUES # Below should be `True` if the current version of `peft` and `transformers` are compatible with From a138d71ec11f2f9c35f98ddae1aebb7f681a1f0b Mon Sep 17 00:00:00 2001 From: YiYi Xu Date: Thu, 23 Oct 2025 22:31:12 -1000 Subject: [PATCH 092/123] HunyuanImage21 (#12333) * add hunyuanimage2.1 --------- Co-authored-by: Sayak Paul --- docs/source/en/_toctree.yml | 8 + .../api/models/autoencoder_kl_hunyuanimage.md | 32 + .../autoencoder_kl_hunyuanimage_refiner.md | 32 + .../api/models/hunyuanimage_transformer_2d.md | 30 + .../source/en/api/pipelines/hunyuanimage21.md | 152 +++ scripts/convert_hunyuan_image_to_diffusers.py | 1044 +++++++++++++++++ src/diffusers/__init__.py | 14 + src/diffusers/guiders/__init__.py | 22 +- .../guiders/adaptive_projected_guidance.py | 54 +- .../adaptive_projected_guidance_mix.py | 284 +++++ src/diffusers/guiders/auto_guidance.py | 14 +- .../guiders/classifier_free_guidance.py | 75 +- .../classifier_free_zero_star_guidance.py | 22 +- .../guiders/frequency_decoupled_guidance.py | 14 +- src/diffusers/guiders/guider_utils.py | 127 +- .../guiders/perturbed_attention_guidance.py | 14 +- src/diffusers/guiders/skip_layer_guidance.py | 14 +- .../guiders/smoothed_energy_guidance.py | 14 +- .../tangential_classifier_free_guidance.py | 16 +- src/diffusers/hooks/_helpers.py | 30 + src/diffusers/models/__init__.py | 6 + src/diffusers/models/autoencoders/__init__.py | 2 + .../autoencoder_kl_hunyuanimage.py | 709 +++++++++++ .../autoencoder_kl_hunyuanimage_refiner.py | 934 +++++++++++++++ src/diffusers/models/transformers/__init__.py | 1 + .../transformers/transformer_hunyuanimage.py | 971 +++++++++++++++ .../modular_pipelines/modular_pipeline.py | 12 +- .../modular_pipelines/qwenimage/denoise.py | 44 +- .../stable_diffusion_xl/denoise.py | 77 +- .../modular_pipelines/wan/denoise.py | 25 +- src/diffusers/pipelines/__init__.py | 2 + .../pipelines/hunyuan_image/__init__.py | 50 + .../hunyuan_image/pipeline_hunyuanimage.py | 866 ++++++++++++++ .../pipeline_hunyuanimage_refiner.py | 752 ++++++++++++ .../hunyuan_image/pipeline_output.py | 21 + .../pipelines/pipeline_loading_utils.py | 1 + src/diffusers/utils/dummy_pt_objects.py | 75 ++ .../dummy_torch_and_transformers_objects.py | 30 + tests/pipelines/hunyuan_image_21/__init__.py | 0 .../hunyuan_image_21/test_hunyuanimage.py | 290 +++++ 40 files changed, 6656 insertions(+), 224 deletions(-) create mode 100644 docs/source/en/api/models/autoencoder_kl_hunyuanimage.md create mode 100644 docs/source/en/api/models/autoencoder_kl_hunyuanimage_refiner.md create mode 100644 docs/source/en/api/models/hunyuanimage_transformer_2d.md create mode 100644 docs/source/en/api/pipelines/hunyuanimage21.md create mode 100644 scripts/convert_hunyuan_image_to_diffusers.py create mode 100644 src/diffusers/guiders/adaptive_projected_guidance_mix.py create mode 100644 src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage.py create mode 100644 src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage_refiner.py create mode 100644 src/diffusers/models/transformers/transformer_hunyuanimage.py create mode 100644 src/diffusers/pipelines/hunyuan_image/__init__.py create mode 100644 src/diffusers/pipelines/hunyuan_image/pipeline_hunyuanimage.py create mode 100644 src/diffusers/pipelines/hunyuan_image/pipeline_hunyuanimage_refiner.py create mode 100644 src/diffusers/pipelines/hunyuan_image/pipeline_output.py create mode 100644 tests/pipelines/hunyuan_image_21/__init__.py create mode 100644 tests/pipelines/hunyuan_image_21/test_hunyuanimage.py diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 1e5c4fe550..540e99a2c6 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -347,6 +347,8 @@ title: HiDreamImageTransformer2DModel - local: api/models/hunyuan_transformer2d title: HunyuanDiT2DModel + - local: api/models/hunyuanimage_transformer_2d + title: HunyuanImageTransformer2DModel - local: api/models/hunyuan_video_transformer_3d title: HunyuanVideoTransformer3DModel - local: api/models/latte_transformer3d @@ -411,6 +413,10 @@ title: AutoencoderKLCogVideoX - local: api/models/autoencoderkl_cosmos title: AutoencoderKLCosmos + - local: api/models/autoencoder_kl_hunyuanimage + title: AutoencoderKLHunyuanImage + - local: api/models/autoencoder_kl_hunyuanimage_refiner + title: AutoencoderKLHunyuanImageRefiner - local: api/models/autoencoder_kl_hunyuan_video title: AutoencoderKLHunyuanVideo - local: api/models/autoencoderkl_ltx_video @@ -620,6 +626,8 @@ title: ConsisID - local: api/pipelines/framepack title: Framepack + - local: api/pipelines/hunyuanimage21 + title: HunyuanImage2.1 - local: api/pipelines/hunyuan_video title: HunyuanVideo - local: api/pipelines/i2vgenxl diff --git a/docs/source/en/api/models/autoencoder_kl_hunyuanimage.md b/docs/source/en/api/models/autoencoder_kl_hunyuanimage.md new file mode 100644 index 0000000000..60dd2b3ab1 --- /dev/null +++ b/docs/source/en/api/models/autoencoder_kl_hunyuanimage.md @@ -0,0 +1,32 @@ + + +# AutoencoderKLHunyuanImage + +The 2D variational autoencoder (VAE) model with KL loss used in [HunyuanImage2.1]. + +The model can be loaded with the following code snippet. + +```python +from diffusers import AutoencoderKLHunyuanImage + +vae = AutoencoderKLHunyuanImage.from_pretrained("hunyuanvideo-community/HunyuanImage-2.1-Diffusers", subfolder="vae", torch_dtype=torch.bfloat16) +``` + +## AutoencoderKLHunyuanImage + +[[autodoc]] AutoencoderKLHunyuanImage + - decode + - all + +## DecoderOutput + +[[autodoc]] models.autoencoders.vae.DecoderOutput diff --git a/docs/source/en/api/models/autoencoder_kl_hunyuanimage_refiner.md b/docs/source/en/api/models/autoencoder_kl_hunyuanimage_refiner.md new file mode 100644 index 0000000000..5e1dd5e2a2 --- /dev/null +++ b/docs/source/en/api/models/autoencoder_kl_hunyuanimage_refiner.md @@ -0,0 +1,32 @@ + + +# AutoencoderKLHunyuanImageRefiner + +The 3D variational autoencoder (VAE) model with KL loss used in [HunyuanImage2.1](https://github.com/Tencent-Hunyuan/HunyuanImage-2.1) for its refiner pipeline. + +The model can be loaded with the following code snippet. + +```python +from diffusers import AutoencoderKLHunyuanImageRefiner + +vae = AutoencoderKLHunyuanImageRefiner.from_pretrained("hunyuanvideo-community/HunyuanImage-2.1-Refiner-Diffusers", subfolder="vae", torch_dtype=torch.bfloat16) +``` + +## AutoencoderKLHunyuanImageRefiner + +[[autodoc]] AutoencoderKLHunyuanImageRefiner + - decode + - all + +## DecoderOutput + +[[autodoc]] models.autoencoders.vae.DecoderOutput diff --git a/docs/source/en/api/models/hunyuanimage_transformer_2d.md b/docs/source/en/api/models/hunyuanimage_transformer_2d.md new file mode 100644 index 0000000000..2dd8e6b52d --- /dev/null +++ b/docs/source/en/api/models/hunyuanimage_transformer_2d.md @@ -0,0 +1,30 @@ + + +# HunyuanImageTransformer2DModel + +A Diffusion Transformer model for [HunyuanImage2.1](https://github.com/Tencent-Hunyuan/HunyuanImage-2.1). + +The model can be loaded with the following code snippet. + +```python +from diffusers import HunyuanImageTransformer2DModel + +transformer = HunyuanImageTransformer2DModel.from_pretrained("hunyuanvideo-community/HunyuanImage-2.1-Diffusers", subfolder="transformer", torch_dtype=torch.bfloat16) +``` + +## HunyuanImageTransformer2DModel + +[[autodoc]] HunyuanImageTransformer2DModel + +## Transformer2DModelOutput + +[[autodoc]] models.modeling_outputs.Transformer2DModelOutput diff --git a/docs/source/en/api/pipelines/hunyuanimage21.md b/docs/source/en/api/pipelines/hunyuanimage21.md new file mode 100644 index 0000000000..f7ba40e237 --- /dev/null +++ b/docs/source/en/api/pipelines/hunyuanimage21.md @@ -0,0 +1,152 @@ + + +# HunyuanImage2.1 + + +HunyuanImage-2.1 is a 17B text-to-image model that is capable of generating 2K (2048 x 2048) resolution images + +HunyuanImage-2.1 comes in the following variants: + +| model type | model id | +|:----------:|:--------:| +| HunyuanImage-2.1 | [hunyuanvideo-community/HunyuanImage-2.1-Diffusers](https://huggingface.co/hunyuanvideo-community/HunyuanImage-2.1-Diffusers) | +| HunyuanImage-2.1-Distilled | [hunyuanvideo-community/HunyuanImage-2.1-Distilled-Diffusers](https://huggingface.co/hunyuanvideo-community/HunyuanImage-2.1-Distilled-Diffusers) | +| HunyuanImage-2.1-Refiner | [hunyuanvideo-community/HunyuanImage-2.1-Refiner-Diffusers](https://huggingface.co/hunyuanvideo-community/HunyuanImage-2.1-Refiner-Diffusers) | + +> [!TIP] +> [Caching](../../optimization/cache) may also speed up inference by storing and reusing intermediate outputs. + +## HunyuanImage-2.1 + +HunyuanImage-2.1 applies [Adaptive Projected Guidance (APG)](https://huggingface.co/papers/2410.02416) combined with Classifier-Free Guidance (CFG) in the denoising loop. `HunyuanImagePipeline` has a `guider` component (read more about [Guider](../modular_diffusers/guiders.md)) and does not take a `guidance_scale` parameter at runtime. To change guider-related parameters, e.g., `guidance_scale`, you can update the `guider` configuration instead. + +```python +import torch +from diffusers import HunyuanImagePipeline + +pipe = HunyuanImagePipeline.from_pretrained( + "hunyuanvideo-community/HunyuanImage-2.1-Diffusers", + torch_dtype=torch.bfloat16 +) +pipe = pipe.to("cuda") +``` + +You can inspect the `guider` object: + +```py +>>> pipe.guider +AdaptiveProjectedMixGuidance { + "_class_name": "AdaptiveProjectedMixGuidance", + "_diffusers_version": "0.36.0.dev0", + "adaptive_projected_guidance_momentum": -0.5, + "adaptive_projected_guidance_rescale": 10.0, + "adaptive_projected_guidance_scale": 10.0, + "adaptive_projected_guidance_start_step": 5, + "enabled": true, + "eta": 0.0, + "guidance_rescale": 0.0, + "guidance_scale": 3.5, + "start": 0.0, + "stop": 1.0, + "use_original_formulation": false +} + +State: + step: None + num_inference_steps: None + timestep: None + count_prepared: 0 + enabled: True + num_conditions: 2 + momentum_buffer: None + is_apg_enabled: False + is_cfg_enabled: True +``` + +To update the guider with a different configuration, use the `new()` method. For example, to generate an image with `guidance_scale=5.0` while keeping all other default guidance parameters: + +```py +import torch +from diffusers import HunyuanImagePipeline + +pipe = HunyuanImagePipeline.from_pretrained( + "hunyuanvideo-community/HunyuanImage-2.1-Diffusers", + torch_dtype=torch.bfloat16 +) +pipe = pipe.to("cuda") + +# Update the guider configuration +pipe.guider = pipe.guider.new(guidance_scale=5.0) + +prompt = ( + "A cute, cartoon-style anthropomorphic penguin plush toy with fluffy fur, standing in a painting studio, " + "wearing a red knitted scarf and a red beret with the word 'Tencent' on it, holding a paintbrush with a " + "focused expression as it paints an oil painting of the Mona Lisa, rendered in a photorealistic photographic style." +) + +image = pipe( + prompt=prompt, + num_inference_steps=50, + height=2048, + width=2048, +).images[0] +image.save("image.png") +``` + + +## HunyuanImage-2.1-Distilled + +use `distilled_guidance_scale` with the guidance-distilled checkpoint, + +```py +import torch +from diffusers import HunyuanImagePipeline +pipe = HunyuanImagePipeline.from_pretrained("hunyuanvideo-community/HunyuanImage-2.1-Distilled-Diffusers", torch_dtype=torch.bfloat16) +pipe = pipe.to("cuda") + +prompt = ( + "A cute, cartoon-style anthropomorphic penguin plush toy with fluffy fur, standing in a painting studio, " + "wearing a red knitted scarf and a red beret with the word 'Tencent' on it, holding a paintbrush with a " + "focused expression as it paints an oil painting of the Mona Lisa, rendered in a photorealistic photographic style." +) + +out = pipe( + prompt, + num_inference_steps=8, + distilled_guidance_scale=3.25, + height=2048, + width=2048, + generator=generator, +).images[0] + +``` + + +## HunyuanImagePipeline + +[[autodoc]] HunyuanImagePipeline + - all + - __call__ + +## HunyuanImageRefinerPipeline + +[[autodoc]] HunyuanImageRefinerPipeline + - all + - __call__ + + +## HunyuanImagePipelineOutput + +[[autodoc]] pipelines.hunyuan_image.pipeline_output.HunyuanImagePipelineOutput \ No newline at end of file diff --git a/scripts/convert_hunyuan_image_to_diffusers.py b/scripts/convert_hunyuan_image_to_diffusers.py new file mode 100644 index 0000000000..c41e934cc3 --- /dev/null +++ b/scripts/convert_hunyuan_image_to_diffusers.py @@ -0,0 +1,1044 @@ +import argparse +import logging + +import torch +from safetensors import safe_open + +from diffusers import AutoencoderKLHunyuanImage, AutoencoderKLHunyuanImageRefiner, HunyuanImageTransformer2DModel + + +logger = logging.getLogger(__name__) # pylint: disable=invalid-name + + +""" +Usage examples +============== + +python scripts/convert_hunyuan_image_to_diffusers.py \ + --model_type hunyuanimage2.1 \ + --transformer_checkpoint_path "/raid/yiyi/HunyuanImage-2.1/ckpts/dit/hunyuanimage2.1.safetensors" \ + --vae_checkpoint_path "HunyuanImage-2.1/ckpts/vae/vae_2_1/pytorch_model.ckpt" \ + --output_path "/raid/yiyi/test-hy21-diffusers" \ + --dtype fp32 + +python scripts/convert_hunyuan_image_to_diffusers.py \ + --model_type hunyuanimage2.1-distilled \ + --transformer_checkpoint_path "/raid/yiyi/HunyuanImage-2.1/ckpts/dit/hunyuanimage2.1-distilled.safetensors" \ + --vae_checkpoint_path "/raid/yiyi/HunyuanImage-2.1/ckpts/vae/vae_2_1/pytorch_model.ckpt" \ + --output_path "/raid/yiyi/test-hy21-distilled-diffusers" \ + --dtype fp32 + + +python scripts/convert_hunyuan_image_to_diffusers.py \ + --model_type hunyuanimage-refiner \ + --transformer_checkpoint_path "/raid/yiyi/HunyuanImage-2.1/ckpts/dit/hunyuanimage-refiner.safetensors" \ + --vae_checkpoint_path "/raid/yiyi/HunyuanImage-2.1/ckpts/vae/vae_refiner/pytorch_model.pt" \ + --output_path "/raid/yiyi/test-hy2-refiner-diffusers" \ + --dtype fp32 +""" + +parser = argparse.ArgumentParser() +parser.add_argument( + "--model_type", type=str, default=None +) # hunyuanimage2.1, hunyuanimage2.1-distilled, hunyuanimage-refiner +parser.add_argument("--transformer_checkpoint_path", default=None, type=str) # ckpts/dit/hunyuanimage2.1.safetensors +parser.add_argument("--vae_checkpoint_path", default=None, type=str) # ckpts/vae/vae_2_1/pytorch_model.ckpt +parser.add_argument("--output_path", type=str) +parser.add_argument("--dtype", type=str, default="fp32") + +args = parser.parse_args() +dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float32 + + +# copied from https://github.com/Tencent-Hunyuan/HunyuanImage-2.1/hyimage/models/hunyuan/modules/hunyuanimage_dit.py#L21 +def convert_hunyuan_dict_for_tensor_parallel(state_dict): + """ + Convert a Hunyuan model state dict to be compatible with tensor parallel architectures. + + Args: + state_dict: Original state dict + + Returns: + new_dict: Converted state dict + """ + new_dict = {} + for k, w in state_dict.items(): + if k.startswith("double_blocks") and "attn_qkv.weight" in k: + hidden_size = w.shape[1] + k1 = k.replace("attn_qkv.weight", "attn_q.weight") + w1 = w[:hidden_size, :] + new_dict[k1] = w1 + k2 = k.replace("attn_qkv.weight", "attn_k.weight") + w2 = w[hidden_size : 2 * hidden_size, :] + new_dict[k2] = w2 + k3 = k.replace("attn_qkv.weight", "attn_v.weight") + w3 = w[-hidden_size:, :] + new_dict[k3] = w3 + elif k.startswith("double_blocks") and "attn_qkv.bias" in k: + hidden_size = w.shape[0] // 3 + k1 = k.replace("attn_qkv.bias", "attn_q.bias") + w1 = w[:hidden_size] + new_dict[k1] = w1 + k2 = k.replace("attn_qkv.bias", "attn_k.bias") + w2 = w[hidden_size : 2 * hidden_size] + new_dict[k2] = w2 + k3 = k.replace("attn_qkv.bias", "attn_v.bias") + w3 = w[-hidden_size:] + new_dict[k3] = w3 + elif k.startswith("single_blocks") and "linear1" in k: + hidden_size = state_dict[k.replace("linear1", "linear2")].shape[0] + k1 = k.replace("linear1", "linear1_q") + w1 = w[:hidden_size] + new_dict[k1] = w1 + k2 = k.replace("linear1", "linear1_k") + w2 = w[hidden_size : 2 * hidden_size] + new_dict[k2] = w2 + k3 = k.replace("linear1", "linear1_v") + w3 = w[2 * hidden_size : 3 * hidden_size] + new_dict[k3] = w3 + k4 = k.replace("linear1", "linear1_mlp") + w4 = w[3 * hidden_size :] + new_dict[k4] = w4 + elif k.startswith("single_blocks") and "linear2" in k: + k1 = k.replace("linear2", "linear2.fc") + new_dict[k1] = w + else: + new_dict[k] = w + return new_dict + + +def load_original_vae_checkpoint(args): + # "ckpts/vae/vae_2_1/pytorch_model.ckpt" + state_dict = torch.load(args.vae_checkpoint_path) + + if "state_dict" in state_dict: + state_dict = state_dict["state_dict"] + vae_state_dict = {} + for k, v in state_dict.items(): + if k.startswith("vae."): + vae_state_dict[k.replace("vae.", "")] = v + + for k, v in vae_state_dict.items(): + if "weight" in k: + if len(v.shape) == 5 and v.shape[2] == 1: + vae_state_dict[k] = v.squeeze(2) + else: + vae_state_dict[k] = v + else: + vae_state_dict[k] = v + return vae_state_dict + + +def load_original_refiner_vae_checkpoint(args): + # "ckpts/vae/vae_refiner/pytorch_model.pt" + state_dict = torch.load(args.vae_checkpoint_path) + + if "state_dict" in state_dict: + state_dict = state_dict["state_dict"] + vae_state_dict = {} + for k, v in state_dict.items(): + if k.startswith("vae."): + vae_state_dict[k.replace("vae.", "")] = v + return vae_state_dict + + +def load_original_transformer_checkpoint(args): + # ckpts/dit/hunyuanimage-refiner.safetensors" + # ckpts/dit/hunyuanimage2.1.safetensors" + state_dict = {} + with safe_open(args.transformer_checkpoint_path, framework="pt", device="cpu") as f: + for key in f.keys(): + state_dict[key] = f.get_tensor(key) + if args.model_type == "hunyuanimage-2.1": + state_dict = convert_hunyuan_dict_for_tensor_parallel(state_dict) + return state_dict + + +def convert_hunyuan_image_transformer_checkpoint_to_diffusers( + original_state_dict, use_byt5=True, guidance_distilled=False, use_meanflow=False +): + converted_state_dict = {} + + # 1. byt5_in -> context_embedder_2 + if use_byt5: + converted_state_dict["context_embedder_2.norm.weight"] = original_state_dict.pop("byt5_in.layernorm.weight") + converted_state_dict["context_embedder_2.norm.bias"] = original_state_dict.pop("byt5_in.layernorm.bias") + converted_state_dict["context_embedder_2.linear_1.weight"] = original_state_dict.pop("byt5_in.fc1.weight") + converted_state_dict["context_embedder_2.linear_1.bias"] = original_state_dict.pop("byt5_in.fc1.bias") + converted_state_dict["context_embedder_2.linear_2.weight"] = original_state_dict.pop("byt5_in.fc2.weight") + converted_state_dict["context_embedder_2.linear_2.bias"] = original_state_dict.pop("byt5_in.fc2.bias") + converted_state_dict["context_embedder_2.linear_3.weight"] = original_state_dict.pop("byt5_in.fc3.weight") + converted_state_dict["context_embedder_2.linear_3.bias"] = original_state_dict.pop("byt5_in.fc3.bias") + + # 2. img_in -> x_embedder + converted_state_dict["x_embedder.proj.weight"] = original_state_dict.pop("img_in.proj.weight") + converted_state_dict["x_embedder.proj.bias"] = original_state_dict.pop("img_in.proj.bias") + + # 3. txt_in -> context_embedder (complex mapping) + # txt_in.input_embedder -> context_embedder.proj_in + converted_state_dict["context_embedder.proj_in.weight"] = original_state_dict.pop("txt_in.input_embedder.weight") + converted_state_dict["context_embedder.proj_in.bias"] = original_state_dict.pop("txt_in.input_embedder.bias") + + # txt_in.t_embedder -> context_embedder.time_text_embed.timestep_embedder + converted_state_dict["context_embedder.time_text_embed.timestep_embedder.linear_1.weight"] = ( + original_state_dict.pop("txt_in.t_embedder.mlp.0.weight") + ) + converted_state_dict["context_embedder.time_text_embed.timestep_embedder.linear_1.bias"] = original_state_dict.pop( + "txt_in.t_embedder.mlp.0.bias" + ) + converted_state_dict["context_embedder.time_text_embed.timestep_embedder.linear_2.weight"] = ( + original_state_dict.pop("txt_in.t_embedder.mlp.2.weight") + ) + converted_state_dict["context_embedder.time_text_embed.timestep_embedder.linear_2.bias"] = original_state_dict.pop( + "txt_in.t_embedder.mlp.2.bias" + ) + + # txt_in.c_embedder -> context_embedder.time_text_embed.text_embedder + converted_state_dict["context_embedder.time_text_embed.text_embedder.linear_1.weight"] = original_state_dict.pop( + "txt_in.c_embedder.linear_1.weight" + ) + converted_state_dict["context_embedder.time_text_embed.text_embedder.linear_1.bias"] = original_state_dict.pop( + "txt_in.c_embedder.linear_1.bias" + ) + converted_state_dict["context_embedder.time_text_embed.text_embedder.linear_2.weight"] = original_state_dict.pop( + "txt_in.c_embedder.linear_2.weight" + ) + converted_state_dict["context_embedder.time_text_embed.text_embedder.linear_2.bias"] = original_state_dict.pop( + "txt_in.c_embedder.linear_2.bias" + ) + + # txt_in.individual_token_refiner -> context_embedder.token_refiner + for i in range(2): # 2 refiner blocks + block_prefix = f"context_embedder.token_refiner.refiner_blocks.{i}." + # norm1 + converted_state_dict[f"{block_prefix}norm1.weight"] = original_state_dict.pop( + f"txt_in.individual_token_refiner.blocks.{i}.norm1.weight" + ) + converted_state_dict[f"{block_prefix}norm1.bias"] = original_state_dict.pop( + f"txt_in.individual_token_refiner.blocks.{i}.norm1.bias" + ) + # norm2 + converted_state_dict[f"{block_prefix}norm2.weight"] = original_state_dict.pop( + f"txt_in.individual_token_refiner.blocks.{i}.norm2.weight" + ) + converted_state_dict[f"{block_prefix}norm2.bias"] = original_state_dict.pop( + f"txt_in.individual_token_refiner.blocks.{i}.norm2.bias" + ) + + # Split QKV + qkv_weight = original_state_dict.pop(f"txt_in.individual_token_refiner.blocks.{i}.self_attn_qkv.weight") + qkv_bias = original_state_dict.pop(f"txt_in.individual_token_refiner.blocks.{i}.self_attn_qkv.bias") + q_weight, k_weight, v_weight = torch.chunk(qkv_weight, 3, dim=0) + q_bias, k_bias, v_bias = torch.chunk(qkv_bias, 3, dim=0) + + converted_state_dict[f"{block_prefix}attn.to_q.weight"] = q_weight + converted_state_dict[f"{block_prefix}attn.to_q.bias"] = q_bias + converted_state_dict[f"{block_prefix}attn.to_k.weight"] = k_weight + converted_state_dict[f"{block_prefix}attn.to_k.bias"] = k_bias + converted_state_dict[f"{block_prefix}attn.to_v.weight"] = v_weight + converted_state_dict[f"{block_prefix}attn.to_v.bias"] = v_bias + + # attn projection + converted_state_dict[f"{block_prefix}attn.to_out.0.weight"] = original_state_dict.pop( + f"txt_in.individual_token_refiner.blocks.{i}.self_attn_proj.weight" + ) + converted_state_dict[f"{block_prefix}attn.to_out.0.bias"] = original_state_dict.pop( + f"txt_in.individual_token_refiner.blocks.{i}.self_attn_proj.bias" + ) + + # MLP + converted_state_dict[f"{block_prefix}ff.net.0.proj.weight"] = original_state_dict.pop( + f"txt_in.individual_token_refiner.blocks.{i}.mlp.fc1.weight" + ) + converted_state_dict[f"{block_prefix}ff.net.0.proj.bias"] = original_state_dict.pop( + f"txt_in.individual_token_refiner.blocks.{i}.mlp.fc1.bias" + ) + converted_state_dict[f"{block_prefix}ff.net.2.weight"] = original_state_dict.pop( + f"txt_in.individual_token_refiner.blocks.{i}.mlp.fc2.weight" + ) + converted_state_dict[f"{block_prefix}ff.net.2.bias"] = original_state_dict.pop( + f"txt_in.individual_token_refiner.blocks.{i}.mlp.fc2.bias" + ) + + # norm_out + converted_state_dict[f"{block_prefix}norm_out.linear.weight"] = original_state_dict.pop( + f"txt_in.individual_token_refiner.blocks.{i}.adaLN_modulation.1.weight" + ) + converted_state_dict[f"{block_prefix}norm_out.linear.bias"] = original_state_dict.pop( + f"txt_in.individual_token_refiner.blocks.{i}.adaLN_modulation.1.bias" + ) + + # 4. time_in -> time_text_embed.timestep_embedder + converted_state_dict["time_guidance_embed.timestep_embedder.linear_1.weight"] = original_state_dict.pop( + "time_in.mlp.0.weight" + ) + converted_state_dict["time_guidance_embed.timestep_embedder.linear_1.bias"] = original_state_dict.pop( + "time_in.mlp.0.bias" + ) + converted_state_dict["time_guidance_embed.timestep_embedder.linear_2.weight"] = original_state_dict.pop( + "time_in.mlp.2.weight" + ) + converted_state_dict["time_guidance_embed.timestep_embedder.linear_2.bias"] = original_state_dict.pop( + "time_in.mlp.2.bias" + ) + + # time_r_in -> time_guidance_embed.timestep_r_embedder + if use_meanflow: + converted_state_dict["time_guidance_embed.timestep_embedder_r.linear_1.weight"] = original_state_dict.pop( + "time_r_in.mlp.0.weight" + ) + converted_state_dict["time_guidance_embed.timestep_embedder_r.linear_1.bias"] = original_state_dict.pop( + "time_r_in.mlp.0.bias" + ) + converted_state_dict["time_guidance_embed.timestep_embedder_r.linear_2.weight"] = original_state_dict.pop( + "time_r_in.mlp.2.weight" + ) + converted_state_dict["time_guidance_embed.timestep_embedder_r.linear_2.bias"] = original_state_dict.pop( + "time_r_in.mlp.2.bias" + ) + + # guidance_in -> time_guidance_embed.guidance_embedder + if guidance_distilled: + converted_state_dict["time_guidance_embed.guidance_embedder.linear_1.weight"] = original_state_dict.pop( + "guidance_in.mlp.0.weight" + ) + converted_state_dict["time_guidance_embed.guidance_embedder.linear_1.bias"] = original_state_dict.pop( + "guidance_in.mlp.0.bias" + ) + converted_state_dict["time_guidance_embed.guidance_embedder.linear_2.weight"] = original_state_dict.pop( + "guidance_in.mlp.2.weight" + ) + converted_state_dict["time_guidance_embed.guidance_embedder.linear_2.bias"] = original_state_dict.pop( + "guidance_in.mlp.2.bias" + ) + + # 5. double_blocks -> transformer_blocks + for i in range(20): # 20 double blocks + block_prefix = f"transformer_blocks.{i}." + + # norm1 (img_mod) + converted_state_dict[f"{block_prefix}norm1.linear.weight"] = original_state_dict.pop( + f"double_blocks.{i}.img_mod.linear.weight" + ) + converted_state_dict[f"{block_prefix}norm1.linear.bias"] = original_state_dict.pop( + f"double_blocks.{i}.img_mod.linear.bias" + ) + + # norm1_context (txt_mod) + converted_state_dict[f"{block_prefix}norm1_context.linear.weight"] = original_state_dict.pop( + f"double_blocks.{i}.txt_mod.linear.weight" + ) + converted_state_dict[f"{block_prefix}norm1_context.linear.bias"] = original_state_dict.pop( + f"double_blocks.{i}.txt_mod.linear.bias" + ) + + # img attention + converted_state_dict[f"{block_prefix}attn.to_q.weight"] = original_state_dict.pop( + f"double_blocks.{i}.img_attn_q.weight" + ) + converted_state_dict[f"{block_prefix}attn.to_q.bias"] = original_state_dict.pop( + f"double_blocks.{i}.img_attn_q.bias" + ) + converted_state_dict[f"{block_prefix}attn.to_k.weight"] = original_state_dict.pop( + f"double_blocks.{i}.img_attn_k.weight" + ) + converted_state_dict[f"{block_prefix}attn.to_k.bias"] = original_state_dict.pop( + f"double_blocks.{i}.img_attn_k.bias" + ) + converted_state_dict[f"{block_prefix}attn.to_v.weight"] = original_state_dict.pop( + f"double_blocks.{i}.img_attn_v.weight" + ) + converted_state_dict[f"{block_prefix}attn.to_v.bias"] = original_state_dict.pop( + f"double_blocks.{i}.img_attn_v.bias" + ) + + # img attention norms + converted_state_dict[f"{block_prefix}attn.norm_q.weight"] = original_state_dict.pop( + f"double_blocks.{i}.img_attn_q_norm.weight" + ) + converted_state_dict[f"{block_prefix}attn.norm_k.weight"] = original_state_dict.pop( + f"double_blocks.{i}.img_attn_k_norm.weight" + ) + + # img attention projection + converted_state_dict[f"{block_prefix}attn.to_out.0.weight"] = original_state_dict.pop( + f"double_blocks.{i}.img_attn_proj.weight" + ) + converted_state_dict[f"{block_prefix}attn.to_out.0.bias"] = original_state_dict.pop( + f"double_blocks.{i}.img_attn_proj.bias" + ) + + # img MLP + converted_state_dict[f"{block_prefix}ff.net.0.proj.weight"] = original_state_dict.pop( + f"double_blocks.{i}.img_mlp.fc1.weight" + ) + converted_state_dict[f"{block_prefix}ff.net.0.proj.bias"] = original_state_dict.pop( + f"double_blocks.{i}.img_mlp.fc1.bias" + ) + converted_state_dict[f"{block_prefix}ff.net.2.weight"] = original_state_dict.pop( + f"double_blocks.{i}.img_mlp.fc2.weight" + ) + converted_state_dict[f"{block_prefix}ff.net.2.bias"] = original_state_dict.pop( + f"double_blocks.{i}.img_mlp.fc2.bias" + ) + + # txt attention (additional projections) + converted_state_dict[f"{block_prefix}attn.add_q_proj.weight"] = original_state_dict.pop( + f"double_blocks.{i}.txt_attn_q.weight" + ) + converted_state_dict[f"{block_prefix}attn.add_q_proj.bias"] = original_state_dict.pop( + f"double_blocks.{i}.txt_attn_q.bias" + ) + converted_state_dict[f"{block_prefix}attn.add_k_proj.weight"] = original_state_dict.pop( + f"double_blocks.{i}.txt_attn_k.weight" + ) + converted_state_dict[f"{block_prefix}attn.add_k_proj.bias"] = original_state_dict.pop( + f"double_blocks.{i}.txt_attn_k.bias" + ) + converted_state_dict[f"{block_prefix}attn.add_v_proj.weight"] = original_state_dict.pop( + f"double_blocks.{i}.txt_attn_v.weight" + ) + converted_state_dict[f"{block_prefix}attn.add_v_proj.bias"] = original_state_dict.pop( + f"double_blocks.{i}.txt_attn_v.bias" + ) + + # txt attention norms + converted_state_dict[f"{block_prefix}attn.norm_added_q.weight"] = original_state_dict.pop( + f"double_blocks.{i}.txt_attn_q_norm.weight" + ) + converted_state_dict[f"{block_prefix}attn.norm_added_k.weight"] = original_state_dict.pop( + f"double_blocks.{i}.txt_attn_k_norm.weight" + ) + + # txt attention projection + converted_state_dict[f"{block_prefix}attn.to_add_out.weight"] = original_state_dict.pop( + f"double_blocks.{i}.txt_attn_proj.weight" + ) + converted_state_dict[f"{block_prefix}attn.to_add_out.bias"] = original_state_dict.pop( + f"double_blocks.{i}.txt_attn_proj.bias" + ) + + # txt MLP (ff_context) + converted_state_dict[f"{block_prefix}ff_context.net.0.proj.weight"] = original_state_dict.pop( + f"double_blocks.{i}.txt_mlp.fc1.weight" + ) + converted_state_dict[f"{block_prefix}ff_context.net.0.proj.bias"] = original_state_dict.pop( + f"double_blocks.{i}.txt_mlp.fc1.bias" + ) + converted_state_dict[f"{block_prefix}ff_context.net.2.weight"] = original_state_dict.pop( + f"double_blocks.{i}.txt_mlp.fc2.weight" + ) + converted_state_dict[f"{block_prefix}ff_context.net.2.bias"] = original_state_dict.pop( + f"double_blocks.{i}.txt_mlp.fc2.bias" + ) + + # 6. single_blocks -> single_transformer_blocks + for i in range(40): # 40 single blocks + block_prefix = f"single_transformer_blocks.{i}." + + # norm + converted_state_dict[f"{block_prefix}norm.linear.weight"] = original_state_dict.pop( + f"single_blocks.{i}.modulation.linear.weight" + ) + converted_state_dict[f"{block_prefix}norm.linear.bias"] = original_state_dict.pop( + f"single_blocks.{i}.modulation.linear.bias" + ) + + # attention Q, K, V + converted_state_dict[f"{block_prefix}attn.to_q.weight"] = original_state_dict.pop( + f"single_blocks.{i}.linear1_q.weight" + ) + converted_state_dict[f"{block_prefix}attn.to_q.bias"] = original_state_dict.pop( + f"single_blocks.{i}.linear1_q.bias" + ) + converted_state_dict[f"{block_prefix}attn.to_k.weight"] = original_state_dict.pop( + f"single_blocks.{i}.linear1_k.weight" + ) + converted_state_dict[f"{block_prefix}attn.to_k.bias"] = original_state_dict.pop( + f"single_blocks.{i}.linear1_k.bias" + ) + converted_state_dict[f"{block_prefix}attn.to_v.weight"] = original_state_dict.pop( + f"single_blocks.{i}.linear1_v.weight" + ) + converted_state_dict[f"{block_prefix}attn.to_v.bias"] = original_state_dict.pop( + f"single_blocks.{i}.linear1_v.bias" + ) + + # attention norms + converted_state_dict[f"{block_prefix}attn.norm_q.weight"] = original_state_dict.pop( + f"single_blocks.{i}.q_norm.weight" + ) + converted_state_dict[f"{block_prefix}attn.norm_k.weight"] = original_state_dict.pop( + f"single_blocks.{i}.k_norm.weight" + ) + + # MLP projection + converted_state_dict[f"{block_prefix}proj_mlp.weight"] = original_state_dict.pop( + f"single_blocks.{i}.linear1_mlp.weight" + ) + converted_state_dict[f"{block_prefix}proj_mlp.bias"] = original_state_dict.pop( + f"single_blocks.{i}.linear1_mlp.bias" + ) + + # output projection + converted_state_dict[f"{block_prefix}proj_out.weight"] = original_state_dict.pop( + f"single_blocks.{i}.linear2.fc.weight" + ) + converted_state_dict[f"{block_prefix}proj_out.bias"] = original_state_dict.pop( + f"single_blocks.{i}.linear2.fc.bias" + ) + + # 7. final_layer -> norm_out + proj_out + converted_state_dict["proj_out.weight"] = original_state_dict.pop("final_layer.linear.weight") + converted_state_dict["proj_out.bias"] = original_state_dict.pop("final_layer.linear.bias") + shift_w, scale_w = original_state_dict.pop("final_layer.adaLN_modulation.1.weight").chunk(2, dim=0) + shift_b, scale_b = original_state_dict.pop("final_layer.adaLN_modulation.1.bias").chunk(2, dim=0) + converted_state_dict["norm_out.linear.weight"] = torch.cat([scale_w, shift_w], dim=0) + converted_state_dict["norm_out.linear.bias"] = torch.cat([scale_b, shift_b], dim=0) + + return converted_state_dict, original_state_dict + + +def convert_hunyuan_image_vae_checkpoint_to_diffusers( + original_state_dict, block_out_channels=[128, 256, 512, 512, 1024, 1024], layers_per_block=2 +): + """Convert original VAE state dict to Diffusers format.""" + converted = {} + + # 1. Encoder + # 1.1 conv_in + converted["encoder.conv_in.weight"] = original_state_dict.pop("encoder.conv_in.weight") + converted["encoder.conv_in.bias"] = original_state_dict.pop("encoder.conv_in.bias") + + # 1.2 down blocks + diffusers_block_idx = 0 + + for block_index in range(len(block_out_channels)): + for resnet_block_index in range(layers_per_block): + orig_prefix = f"encoder.down.{block_index}.block.{resnet_block_index}" + diff_prefix = f"encoder.down_blocks.{diffusers_block_idx}" + + # resnet blocks + converted[f"{diff_prefix}.norm1.weight"] = original_state_dict.pop(f"{orig_prefix}.norm1.weight") + converted[f"{diff_prefix}.norm1.bias"] = original_state_dict.pop(f"{orig_prefix}.norm1.bias") + converted[f"{diff_prefix}.conv1.weight"] = original_state_dict.pop(f"{orig_prefix}.conv1.weight") + converted[f"{diff_prefix}.conv1.bias"] = original_state_dict.pop(f"{orig_prefix}.conv1.bias") + converted[f"{diff_prefix}.norm2.weight"] = original_state_dict.pop(f"{orig_prefix}.norm2.weight") + converted[f"{diff_prefix}.norm2.bias"] = original_state_dict.pop(f"{orig_prefix}.norm2.bias") + converted[f"{diff_prefix}.conv2.weight"] = original_state_dict.pop(f"{orig_prefix}.conv2.weight") + converted[f"{diff_prefix}.conv2.bias"] = original_state_dict.pop(f"{orig_prefix}.conv2.bias") + + diffusers_block_idx += 1 + + # downsample blocks + if f"encoder.down.{block_index}.downsample.conv.weight" in original_state_dict: + converted[f"encoder.down_blocks.{diffusers_block_idx}.conv.weight"] = original_state_dict.pop( + f"encoder.down.{block_index}.downsample.conv.weight" + ) + converted[f"encoder.down_blocks.{diffusers_block_idx}.conv.bias"] = original_state_dict.pop( + f"encoder.down.{block_index}.downsample.conv.bias" + ) + diffusers_block_idx += 1 + + # 1.3 mid block + converted["encoder.mid_block.resnets.0.norm1.weight"] = original_state_dict.pop("encoder.mid.block_1.norm1.weight") + converted["encoder.mid_block.resnets.0.norm1.bias"] = original_state_dict.pop("encoder.mid.block_1.norm1.bias") + converted["encoder.mid_block.resnets.0.conv1.weight"] = original_state_dict.pop("encoder.mid.block_1.conv1.weight") + converted["encoder.mid_block.resnets.0.conv1.bias"] = original_state_dict.pop("encoder.mid.block_1.conv1.bias") + converted["encoder.mid_block.resnets.0.norm2.weight"] = original_state_dict.pop("encoder.mid.block_1.norm2.weight") + converted["encoder.mid_block.resnets.0.norm2.bias"] = original_state_dict.pop("encoder.mid.block_1.norm2.bias") + converted["encoder.mid_block.resnets.0.conv2.weight"] = original_state_dict.pop("encoder.mid.block_1.conv2.weight") + converted["encoder.mid_block.resnets.0.conv2.bias"] = original_state_dict.pop("encoder.mid.block_1.conv2.bias") + + converted["encoder.mid_block.resnets.1.norm1.weight"] = original_state_dict.pop("encoder.mid.block_2.norm1.weight") + converted["encoder.mid_block.resnets.1.norm1.bias"] = original_state_dict.pop("encoder.mid.block_2.norm1.bias") + converted["encoder.mid_block.resnets.1.conv1.weight"] = original_state_dict.pop("encoder.mid.block_2.conv1.weight") + converted["encoder.mid_block.resnets.1.conv1.bias"] = original_state_dict.pop("encoder.mid.block_2.conv1.bias") + converted["encoder.mid_block.resnets.1.norm2.weight"] = original_state_dict.pop("encoder.mid.block_2.norm2.weight") + converted["encoder.mid_block.resnets.1.norm2.bias"] = original_state_dict.pop("encoder.mid.block_2.norm2.bias") + converted["encoder.mid_block.resnets.1.conv2.weight"] = original_state_dict.pop("encoder.mid.block_2.conv2.weight") + converted["encoder.mid_block.resnets.1.conv2.bias"] = original_state_dict.pop("encoder.mid.block_2.conv2.bias") + + converted["encoder.mid_block.attentions.0.norm.weight"] = original_state_dict.pop("encoder.mid.attn_1.norm.weight") + converted["encoder.mid_block.attentions.0.norm.bias"] = original_state_dict.pop("encoder.mid.attn_1.norm.bias") + converted["encoder.mid_block.attentions.0.to_q.weight"] = original_state_dict.pop("encoder.mid.attn_1.q.weight") + converted["encoder.mid_block.attentions.0.to_q.bias"] = original_state_dict.pop("encoder.mid.attn_1.q.bias") + converted["encoder.mid_block.attentions.0.to_k.weight"] = original_state_dict.pop("encoder.mid.attn_1.k.weight") + converted["encoder.mid_block.attentions.0.to_k.bias"] = original_state_dict.pop("encoder.mid.attn_1.k.bias") + converted["encoder.mid_block.attentions.0.to_v.weight"] = original_state_dict.pop("encoder.mid.attn_1.v.weight") + converted["encoder.mid_block.attentions.0.to_v.bias"] = original_state_dict.pop("encoder.mid.attn_1.v.bias") + converted["encoder.mid_block.attentions.0.proj.weight"] = original_state_dict.pop( + "encoder.mid.attn_1.proj_out.weight" + ) + converted["encoder.mid_block.attentions.0.proj.bias"] = original_state_dict.pop("encoder.mid.attn_1.proj_out.bias") + + # 1.4 encoder output + converted["encoder.norm_out.weight"] = original_state_dict.pop("encoder.norm_out.weight") + converted["encoder.norm_out.bias"] = original_state_dict.pop("encoder.norm_out.bias") + converted["encoder.conv_out.weight"] = original_state_dict.pop("encoder.conv_out.weight") + converted["encoder.conv_out.bias"] = original_state_dict.pop("encoder.conv_out.bias") + + # 2. Decoder + # 2.1 conv_in + converted["decoder.conv_in.weight"] = original_state_dict.pop("decoder.conv_in.weight") + converted["decoder.conv_in.bias"] = original_state_dict.pop("decoder.conv_in.bias") + + # 2.2 mid block + converted["decoder.mid_block.resnets.0.norm1.weight"] = original_state_dict.pop("decoder.mid.block_1.norm1.weight") + converted["decoder.mid_block.resnets.0.norm1.bias"] = original_state_dict.pop("decoder.mid.block_1.norm1.bias") + converted["decoder.mid_block.resnets.0.conv1.weight"] = original_state_dict.pop("decoder.mid.block_1.conv1.weight") + converted["decoder.mid_block.resnets.0.conv1.bias"] = original_state_dict.pop("decoder.mid.block_1.conv1.bias") + converted["decoder.mid_block.resnets.0.norm2.weight"] = original_state_dict.pop("decoder.mid.block_1.norm2.weight") + converted["decoder.mid_block.resnets.0.norm2.bias"] = original_state_dict.pop("decoder.mid.block_1.norm2.bias") + converted["decoder.mid_block.resnets.0.conv2.weight"] = original_state_dict.pop("decoder.mid.block_1.conv2.weight") + converted["decoder.mid_block.resnets.0.conv2.bias"] = original_state_dict.pop("decoder.mid.block_1.conv2.bias") + + converted["decoder.mid_block.resnets.1.norm1.weight"] = original_state_dict.pop("decoder.mid.block_2.norm1.weight") + converted["decoder.mid_block.resnets.1.norm1.bias"] = original_state_dict.pop("decoder.mid.block_2.norm1.bias") + converted["decoder.mid_block.resnets.1.conv1.weight"] = original_state_dict.pop("decoder.mid.block_2.conv1.weight") + converted["decoder.mid_block.resnets.1.conv1.bias"] = original_state_dict.pop("decoder.mid.block_2.conv1.bias") + converted["decoder.mid_block.resnets.1.norm2.weight"] = original_state_dict.pop("decoder.mid.block_2.norm2.weight") + converted["decoder.mid_block.resnets.1.norm2.bias"] = original_state_dict.pop("decoder.mid.block_2.norm2.bias") + converted["decoder.mid_block.resnets.1.conv2.weight"] = original_state_dict.pop("decoder.mid.block_2.conv2.weight") + converted["decoder.mid_block.resnets.1.conv2.bias"] = original_state_dict.pop("decoder.mid.block_2.conv2.bias") + + converted["decoder.mid_block.attentions.0.norm.weight"] = original_state_dict.pop("decoder.mid.attn_1.norm.weight") + converted["decoder.mid_block.attentions.0.norm.bias"] = original_state_dict.pop("decoder.mid.attn_1.norm.bias") + converted["decoder.mid_block.attentions.0.to_q.weight"] = original_state_dict.pop("decoder.mid.attn_1.q.weight") + converted["decoder.mid_block.attentions.0.to_q.bias"] = original_state_dict.pop("decoder.mid.attn_1.q.bias") + converted["decoder.mid_block.attentions.0.to_k.weight"] = original_state_dict.pop("decoder.mid.attn_1.k.weight") + converted["decoder.mid_block.attentions.0.to_k.bias"] = original_state_dict.pop("decoder.mid.attn_1.k.bias") + converted["decoder.mid_block.attentions.0.to_v.weight"] = original_state_dict.pop("decoder.mid.attn_1.v.weight") + converted["decoder.mid_block.attentions.0.to_v.bias"] = original_state_dict.pop("decoder.mid.attn_1.v.bias") + converted["decoder.mid_block.attentions.0.proj.weight"] = original_state_dict.pop( + "decoder.mid.attn_1.proj_out.weight" + ) + converted["decoder.mid_block.attentions.0.proj.bias"] = original_state_dict.pop("decoder.mid.attn_1.proj_out.bias") + + # 2.3 up blocks + diffusers_block_idx = 0 + for up_block_index in range(len(block_out_channels)): + # resnet blocks + for resnet_block_index in range(layers_per_block + 1): + orig_prefix = f"decoder.up.{up_block_index}.block.{resnet_block_index}" + diff_prefix = f"decoder.up_blocks.{diffusers_block_idx}" + + converted[f"{diff_prefix}.norm1.weight"] = original_state_dict.pop(f"{orig_prefix}.norm1.weight") + converted[f"{diff_prefix}.norm1.bias"] = original_state_dict.pop(f"{orig_prefix}.norm1.bias") + converted[f"{diff_prefix}.conv1.weight"] = original_state_dict.pop(f"{orig_prefix}.conv1.weight") + converted[f"{diff_prefix}.conv1.bias"] = original_state_dict.pop(f"{orig_prefix}.conv1.bias") + converted[f"{diff_prefix}.norm2.weight"] = original_state_dict.pop(f"{orig_prefix}.norm2.weight") + converted[f"{diff_prefix}.norm2.bias"] = original_state_dict.pop(f"{orig_prefix}.norm2.bias") + converted[f"{diff_prefix}.conv2.weight"] = original_state_dict.pop(f"{orig_prefix}.conv2.weight") + converted[f"{diff_prefix}.conv2.bias"] = original_state_dict.pop(f"{orig_prefix}.conv2.bias") + + diffusers_block_idx += 1 + + # upsample blocks + if f"decoder.up.{up_block_index}.upsample.conv.weight" in original_state_dict: + converted[f"decoder.up_blocks.{diffusers_block_idx}.conv.weight"] = original_state_dict.pop( + f"decoder.up.{up_block_index}.upsample.conv.weight" + ) + converted[f"decoder.up_blocks.{diffusers_block_idx}.conv.bias"] = original_state_dict.pop( + f"decoder.up.{up_block_index}.upsample.conv.bias" + ) + diffusers_block_idx += 1 + + # 2.4 decoder output + converted["decoder.norm_out.weight"] = original_state_dict.pop("decoder.norm_out.weight") + converted["decoder.norm_out.bias"] = original_state_dict.pop("decoder.norm_out.bias") + converted["decoder.conv_out.weight"] = original_state_dict.pop("decoder.conv_out.weight") + converted["decoder.conv_out.bias"] = original_state_dict.pop("decoder.conv_out.bias") + + return converted, original_state_dict + + +def convert_hunyuan_image_refiner_vae_checkpoint_to_diffusers( + original_state_dict, block_out_channels=[128, 256, 512, 1024, 1024], layers_per_block=2 +): + converted = {} + + # 1. Encoder + # 1.1 conv_in + converted["encoder.conv_in.conv.weight"] = original_state_dict.pop("encoder.conv_in.conv.weight") + converted["encoder.conv_in.conv.bias"] = original_state_dict.pop("encoder.conv_in.conv.bias") + + # 1.2 Down blocks + for down_block_index in range(len(block_out_channels)): # 0 to 4 + # ResNet blocks + for resnet_block_index in range(layers_per_block): # 0 to 1 + converted[f"encoder.down_blocks.{down_block_index}.resnets.{resnet_block_index}.norm1.gamma"] = ( + original_state_dict.pop(f"encoder.down.{down_block_index}.block.{resnet_block_index}.norm1.gamma") + ) + converted[f"encoder.down_blocks.{down_block_index}.resnets.{resnet_block_index}.conv1.conv.weight"] = ( + original_state_dict.pop( + f"encoder.down.{down_block_index}.block.{resnet_block_index}.conv1.conv.weight" + ) + ) + converted[f"encoder.down_blocks.{down_block_index}.resnets.{resnet_block_index}.conv1.conv.bias"] = ( + original_state_dict.pop(f"encoder.down.{down_block_index}.block.{resnet_block_index}.conv1.conv.bias") + ) + converted[f"encoder.down_blocks.{down_block_index}.resnets.{resnet_block_index}.norm2.gamma"] = ( + original_state_dict.pop(f"encoder.down.{down_block_index}.block.{resnet_block_index}.norm2.gamma") + ) + converted[f"encoder.down_blocks.{down_block_index}.resnets.{resnet_block_index}.conv2.conv.weight"] = ( + original_state_dict.pop( + f"encoder.down.{down_block_index}.block.{resnet_block_index}.conv2.conv.weight" + ) + ) + converted[f"encoder.down_blocks.{down_block_index}.resnets.{resnet_block_index}.conv2.conv.bias"] = ( + original_state_dict.pop(f"encoder.down.{down_block_index}.block.{resnet_block_index}.conv2.conv.bias") + ) + + # Downsample (if exists) + if f"encoder.down.{down_block_index}.downsample.conv.conv.weight" in original_state_dict: + converted[f"encoder.down_blocks.{down_block_index}.downsamplers.0.conv.conv.weight"] = ( + original_state_dict.pop(f"encoder.down.{down_block_index}.downsample.conv.conv.weight") + ) + converted[f"encoder.down_blocks.{down_block_index}.downsamplers.0.conv.conv.bias"] = ( + original_state_dict.pop(f"encoder.down.{down_block_index}.downsample.conv.conv.bias") + ) + + # 1.3 Mid block + converted["encoder.mid_block.resnets.0.norm1.gamma"] = original_state_dict.pop("encoder.mid.block_1.norm1.gamma") + converted["encoder.mid_block.resnets.0.conv1.conv.weight"] = original_state_dict.pop( + "encoder.mid.block_1.conv1.conv.weight" + ) + converted["encoder.mid_block.resnets.0.conv1.conv.bias"] = original_state_dict.pop( + "encoder.mid.block_1.conv1.conv.bias" + ) + converted["encoder.mid_block.resnets.0.norm2.gamma"] = original_state_dict.pop("encoder.mid.block_1.norm2.gamma") + converted["encoder.mid_block.resnets.0.conv2.conv.weight"] = original_state_dict.pop( + "encoder.mid.block_1.conv2.conv.weight" + ) + converted["encoder.mid_block.resnets.0.conv2.conv.bias"] = original_state_dict.pop( + "encoder.mid.block_1.conv2.conv.bias" + ) + + converted["encoder.mid_block.resnets.1.norm1.gamma"] = original_state_dict.pop("encoder.mid.block_2.norm1.gamma") + converted["encoder.mid_block.resnets.1.conv1.conv.weight"] = original_state_dict.pop( + "encoder.mid.block_2.conv1.conv.weight" + ) + converted["encoder.mid_block.resnets.1.conv1.conv.bias"] = original_state_dict.pop( + "encoder.mid.block_2.conv1.conv.bias" + ) + converted["encoder.mid_block.resnets.1.norm2.gamma"] = original_state_dict.pop("encoder.mid.block_2.norm2.gamma") + converted["encoder.mid_block.resnets.1.conv2.conv.weight"] = original_state_dict.pop( + "encoder.mid.block_2.conv2.conv.weight" + ) + converted["encoder.mid_block.resnets.1.conv2.conv.bias"] = original_state_dict.pop( + "encoder.mid.block_2.conv2.conv.bias" + ) + + # Attention block + converted["encoder.mid_block.attentions.0.norm.gamma"] = original_state_dict.pop("encoder.mid.attn_1.norm.gamma") + converted["encoder.mid_block.attentions.0.to_q.weight"] = original_state_dict.pop("encoder.mid.attn_1.q.weight") + converted["encoder.mid_block.attentions.0.to_q.bias"] = original_state_dict.pop("encoder.mid.attn_1.q.bias") + converted["encoder.mid_block.attentions.0.to_k.weight"] = original_state_dict.pop("encoder.mid.attn_1.k.weight") + converted["encoder.mid_block.attentions.0.to_k.bias"] = original_state_dict.pop("encoder.mid.attn_1.k.bias") + converted["encoder.mid_block.attentions.0.to_v.weight"] = original_state_dict.pop("encoder.mid.attn_1.v.weight") + converted["encoder.mid_block.attentions.0.to_v.bias"] = original_state_dict.pop("encoder.mid.attn_1.v.bias") + converted["encoder.mid_block.attentions.0.proj_out.weight"] = original_state_dict.pop( + "encoder.mid.attn_1.proj_out.weight" + ) + converted["encoder.mid_block.attentions.0.proj_out.bias"] = original_state_dict.pop( + "encoder.mid.attn_1.proj_out.bias" + ) + + # 1.4 Encoder output + converted["encoder.norm_out.gamma"] = original_state_dict.pop("encoder.norm_out.gamma") + converted["encoder.conv_out.conv.weight"] = original_state_dict.pop("encoder.conv_out.conv.weight") + converted["encoder.conv_out.conv.bias"] = original_state_dict.pop("encoder.conv_out.conv.bias") + + # 2. Decoder + # 2.1 conv_in + converted["decoder.conv_in.conv.weight"] = original_state_dict.pop("decoder.conv_in.conv.weight") + converted["decoder.conv_in.conv.bias"] = original_state_dict.pop("decoder.conv_in.conv.bias") + + # 2.2 Mid block + converted["decoder.mid_block.resnets.0.norm1.gamma"] = original_state_dict.pop("decoder.mid.block_1.norm1.gamma") + converted["decoder.mid_block.resnets.0.conv1.conv.weight"] = original_state_dict.pop( + "decoder.mid.block_1.conv1.conv.weight" + ) + converted["decoder.mid_block.resnets.0.conv1.conv.bias"] = original_state_dict.pop( + "decoder.mid.block_1.conv1.conv.bias" + ) + converted["decoder.mid_block.resnets.0.norm2.gamma"] = original_state_dict.pop("decoder.mid.block_1.norm2.gamma") + converted["decoder.mid_block.resnets.0.conv2.conv.weight"] = original_state_dict.pop( + "decoder.mid.block_1.conv2.conv.weight" + ) + converted["decoder.mid_block.resnets.0.conv2.conv.bias"] = original_state_dict.pop( + "decoder.mid.block_1.conv2.conv.bias" + ) + + converted["decoder.mid_block.resnets.1.norm1.gamma"] = original_state_dict.pop("decoder.mid.block_2.norm1.gamma") + converted["decoder.mid_block.resnets.1.conv1.conv.weight"] = original_state_dict.pop( + "decoder.mid.block_2.conv1.conv.weight" + ) + converted["decoder.mid_block.resnets.1.conv1.conv.bias"] = original_state_dict.pop( + "decoder.mid.block_2.conv1.conv.bias" + ) + converted["decoder.mid_block.resnets.1.norm2.gamma"] = original_state_dict.pop("decoder.mid.block_2.norm2.gamma") + converted["decoder.mid_block.resnets.1.conv2.conv.weight"] = original_state_dict.pop( + "decoder.mid.block_2.conv2.conv.weight" + ) + converted["decoder.mid_block.resnets.1.conv2.conv.bias"] = original_state_dict.pop( + "decoder.mid.block_2.conv2.conv.bias" + ) + + # Decoder attention block + converted["decoder.mid_block.attentions.0.norm.gamma"] = original_state_dict.pop("decoder.mid.attn_1.norm.gamma") + converted["decoder.mid_block.attentions.0.to_q.weight"] = original_state_dict.pop("decoder.mid.attn_1.q.weight") + converted["decoder.mid_block.attentions.0.to_q.bias"] = original_state_dict.pop("decoder.mid.attn_1.q.bias") + converted["decoder.mid_block.attentions.0.to_k.weight"] = original_state_dict.pop("decoder.mid.attn_1.k.weight") + converted["decoder.mid_block.attentions.0.to_k.bias"] = original_state_dict.pop("decoder.mid.attn_1.k.bias") + converted["decoder.mid_block.attentions.0.to_v.weight"] = original_state_dict.pop("decoder.mid.attn_1.v.weight") + converted["decoder.mid_block.attentions.0.to_v.bias"] = original_state_dict.pop("decoder.mid.attn_1.v.bias") + converted["decoder.mid_block.attentions.0.proj_out.weight"] = original_state_dict.pop( + "decoder.mid.attn_1.proj_out.weight" + ) + converted["decoder.mid_block.attentions.0.proj_out.bias"] = original_state_dict.pop( + "decoder.mid.attn_1.proj_out.bias" + ) + + # 2.3 Up blocks + for up_block_index in range(len(block_out_channels)): # 0 to 5 + # ResNet blocks + for resnet_block_index in range(layers_per_block + 1): # 0 to 2 (decoder has 3 resnets per level) + converted[f"decoder.up_blocks.{up_block_index}.resnets.{resnet_block_index}.norm1.gamma"] = ( + original_state_dict.pop(f"decoder.up.{up_block_index}.block.{resnet_block_index}.norm1.gamma") + ) + converted[f"decoder.up_blocks.{up_block_index}.resnets.{resnet_block_index}.conv1.conv.weight"] = ( + original_state_dict.pop(f"decoder.up.{up_block_index}.block.{resnet_block_index}.conv1.conv.weight") + ) + converted[f"decoder.up_blocks.{up_block_index}.resnets.{resnet_block_index}.conv1.conv.bias"] = ( + original_state_dict.pop(f"decoder.up.{up_block_index}.block.{resnet_block_index}.conv1.conv.bias") + ) + converted[f"decoder.up_blocks.{up_block_index}.resnets.{resnet_block_index}.norm2.gamma"] = ( + original_state_dict.pop(f"decoder.up.{up_block_index}.block.{resnet_block_index}.norm2.gamma") + ) + converted[f"decoder.up_blocks.{up_block_index}.resnets.{resnet_block_index}.conv2.conv.weight"] = ( + original_state_dict.pop(f"decoder.up.{up_block_index}.block.{resnet_block_index}.conv2.conv.weight") + ) + converted[f"decoder.up_blocks.{up_block_index}.resnets.{resnet_block_index}.conv2.conv.bias"] = ( + original_state_dict.pop(f"decoder.up.{up_block_index}.block.{resnet_block_index}.conv2.conv.bias") + ) + + # Upsample (if exists) + if f"decoder.up.{up_block_index}.upsample.conv.conv.weight" in original_state_dict: + converted[f"decoder.up_blocks.{up_block_index}.upsamplers.0.conv.conv.weight"] = original_state_dict.pop( + f"decoder.up.{up_block_index}.upsample.conv.conv.weight" + ) + converted[f"decoder.up_blocks.{up_block_index}.upsamplers.0.conv.conv.bias"] = original_state_dict.pop( + f"decoder.up.{up_block_index}.upsample.conv.conv.bias" + ) + + # 2.4 Decoder output + converted["decoder.norm_out.gamma"] = original_state_dict.pop("decoder.norm_out.gamma") + converted["decoder.conv_out.conv.weight"] = original_state_dict.pop("decoder.conv_out.conv.weight") + converted["decoder.conv_out.conv.bias"] = original_state_dict.pop("decoder.conv_out.conv.bias") + + return converted, original_state_dict + + +def main(args): + if args.model_type == "hunyuanimage2.1": + original_transformer_state_dict = load_original_transformer_checkpoint(args) + original_vae_state_dict = load_original_vae_checkpoint(args) + + transformer_config = { + "in_channels": 64, + "out_channels": 64, + "num_attention_heads": 28, + "attention_head_dim": 128, + "num_layers": 20, + "num_single_layers": 40, + "num_refiner_layers": 2, + "patch_size": (1, 1), + "qk_norm": "rms_norm", + "guidance_embeds": False, + "text_embed_dim": 3584, + "text_embed_2_dim": 1472, + "rope_theta": 256.0, + "rope_axes_dim": (64, 64), + } + + converted_transformer_state_dict, original_transformer_state_dict = ( + convert_hunyuan_image_transformer_checkpoint_to_diffusers( + original_transformer_state_dict, use_byt5=True, guidance_distilled=False + ) + ) + + if original_transformer_state_dict: + logger.warning( + f"Unused {len(original_transformer_state_dict)} original keys for transformer: {list(original_transformer_state_dict.keys())}" + ) + + transformer = HunyuanImageTransformer2DModel(**transformer_config) + missing_keys, unexpected_key = transformer.load_state_dict(converted_transformer_state_dict, strict=True) + + if missing_keys: + logger.warning(f"Missing keys for transformer: {missing_keys}") + if unexpected_key: + logger.warning(f"Unexpected keys for transformer: {unexpected_key}") + + transformer.to(dtype).save_pretrained(f"{args.output_path}/transformer") + + vae_config_diffusers = { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 64, + "block_out_channels": [128, 256, 512, 512, 1024, 1024], + "layers_per_block": 2, + "spatial_compression_ratio": 32, + "sample_size": 384, + "scaling_factor": 0.75289, + "downsample_match_channel": True, + "upsample_match_channel": True, + } + converted_vae_state_dict, original_vae_state_dict = convert_hunyuan_image_vae_checkpoint_to_diffusers( + original_vae_state_dict, block_out_channels=[128, 256, 512, 512, 1024, 1024], layers_per_block=2 + ) + if original_vae_state_dict: + logger.warning( + f"Unused {len(original_vae_state_dict)} original keys for vae: {list(original_vae_state_dict.keys())}" + ) + + vae = AutoencoderKLHunyuanImage(**vae_config_diffusers) + missing_keys, unexpected_key = vae.load_state_dict(converted_vae_state_dict, strict=True) + + if missing_keys: + logger.warning(f"Missing keys for vae: {missing_keys}") + if unexpected_key: + logger.warning(f"Unexpected keys for vae: {unexpected_key}") + + vae.to(dtype).save_pretrained(f"{args.output_path}/vae") + + elif args.model_type == "hunyuanimage2.1-distilled": + original_transformer_state_dict = load_original_transformer_checkpoint(args) + original_vae_state_dict = load_original_vae_checkpoint(args) + + transformer_config = { + "in_channels": 64, + "out_channels": 64, + "num_attention_heads": 28, + "attention_head_dim": 128, + "num_layers": 20, + "num_single_layers": 40, + "num_refiner_layers": 2, + "patch_size": (1, 1), + "qk_norm": "rms_norm", + "guidance_embeds": True, + "text_embed_dim": 3584, + "text_embed_2_dim": 1472, + "rope_theta": 256.0, + "rope_axes_dim": (64, 64), + "use_meanflow": True, + } + + converted_transformer_state_dict, original_transformer_state_dict = ( + convert_hunyuan_image_transformer_checkpoint_to_diffusers( + original_transformer_state_dict, use_byt5=True, guidance_distilled=True, use_meanflow=True + ) + ) + + if original_transformer_state_dict: + logger.warning( + f"Unused {len(original_transformer_state_dict)} original keys for transformer: {list(original_transformer_state_dict.keys())}" + ) + + transformer = HunyuanImageTransformer2DModel(**transformer_config) + missing_keys, unexpected_key = transformer.load_state_dict(converted_transformer_state_dict, strict=True) + + if missing_keys: + logger.warning(f"Missing keys for transformer: {missing_keys}") + if unexpected_key: + logger.warning(f"Unexpected keys for transformer: {unexpected_key}") + + transformer.to(dtype).save_pretrained(f"{args.output_path}/transformer") + + vae_config_diffusers = { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 64, + "block_out_channels": [128, 256, 512, 512, 1024, 1024], + "layers_per_block": 2, + "spatial_compression_ratio": 32, + "sample_size": 384, + "scaling_factor": 0.75289, + "downsample_match_channel": True, + "upsample_match_channel": True, + } + converted_vae_state_dict, original_vae_state_dict = convert_hunyuan_image_vae_checkpoint_to_diffusers( + original_vae_state_dict, block_out_channels=[128, 256, 512, 512, 1024, 1024], layers_per_block=2 + ) + if original_vae_state_dict: + logger.warning( + f"Unused {len(original_vae_state_dict)} original keys for vae: {list(original_vae_state_dict.keys())}" + ) + + vae = AutoencoderKLHunyuanImage(**vae_config_diffusers) + missing_keys, unexpected_key = vae.load_state_dict(converted_vae_state_dict, strict=True) + + if missing_keys: + logger.warning(f"Missing keys for vae: {missing_keys}") + if unexpected_key: + logger.warning(f"Unexpected keys for vae: {unexpected_key}") + + vae.to(dtype).save_pretrained(f"{args.output_path}/vae") + + elif args.model_type == "hunyuanimage-refiner": + original_transformer_state_dict = load_original_transformer_checkpoint(args) + original_vae_state_dict = load_original_refiner_vae_checkpoint(args) + + transformer_config = { + "in_channels": 128, + "out_channels": 64, + "num_layers": 20, + "num_single_layers": 40, + "rope_axes_dim": [16, 56, 56], + "num_attention_heads": 26, + "attention_head_dim": 128, + "mlp_ratio": 4, + "patch_size": (1, 1, 1), + "text_embed_dim": 3584, + "guidance_embeds": True, + } + converted_transformer_state_dict, original_transformer_state_dict = ( + convert_hunyuan_image_transformer_checkpoint_to_diffusers( + original_transformer_state_dict, use_byt5=False, guidance_distilled=True + ) + ) + if original_transformer_state_dict: + logger.warning( + f"Unused {len(original_transformer_state_dict)} original keys for transformer: {list(original_transformer_state_dict.keys())}" + ) + + transformer = HunyuanImageTransformer2DModel(**transformer_config) + missing_keys, unexpected_key = transformer.load_state_dict(converted_transformer_state_dict, strict=True) + if missing_keys: + logger.warning(f"Missing keys for transformer: {missing_keys}") + if unexpected_key: + logger.warning(f"Unexpected keys for transformer: {unexpected_key}") + + transformer.to(dtype).save_pretrained(f"{args.output_path}/transformer") + + vae = AutoencoderKLHunyuanImageRefiner() + + converted_vae_state_dict, original_vae_state_dict = convert_hunyuan_image_refiner_vae_checkpoint_to_diffusers( + original_vae_state_dict + ) + if original_vae_state_dict: + logger.warning( + f"Unused {len(original_vae_state_dict)} original keys for vae: {list(original_vae_state_dict.keys())}" + ) + + missing_keys, unexpected_key = vae.load_state_dict(converted_vae_state_dict, strict=True) + logger.warning(f"Missing keys for vae: {missing_keys}") + logger.warning(f"Unexpected keys for vae: {unexpected_key}") + + vae.to(dtype).save_pretrained(f"{args.output_path}/vae") + + +if __name__ == "__main__": + main(args) diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index d9c6ba9bbf..2a8b589846 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -149,7 +149,9 @@ else: _import_structure["guiders"].extend( [ "AdaptiveProjectedGuidance", + "AdaptiveProjectedMixGuidance", "AutoGuidance", + "BaseGuidance", "ClassifierFreeGuidance", "ClassifierFreeZeroStarGuidance", "FrequencyDecoupledGuidance", @@ -184,6 +186,8 @@ else: "AutoencoderKLAllegro", "AutoencoderKLCogVideoX", "AutoencoderKLCosmos", + "AutoencoderKLHunyuanImage", + "AutoencoderKLHunyuanImageRefiner", "AutoencoderKLHunyuanVideo", "AutoencoderKLLTXVideo", "AutoencoderKLMagvit", @@ -216,6 +220,7 @@ else: "HunyuanDiT2DControlNetModel", "HunyuanDiT2DModel", "HunyuanDiT2DMultiControlNetModel", + "HunyuanImageTransformer2DModel", "HunyuanVideoFramepackTransformer3DModel", "HunyuanVideoTransformer3DModel", "I2VGenXLUNet", @@ -462,6 +467,8 @@ else: "HunyuanDiTControlNetPipeline", "HunyuanDiTPAGPipeline", "HunyuanDiTPipeline", + "HunyuanImagePipeline", + "HunyuanImageRefinerPipeline", "HunyuanSkyreelsImageToVideoPipeline", "HunyuanVideoFramepackPipeline", "HunyuanVideoImageToVideoPipeline", @@ -849,7 +856,9 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: else: from .guiders import ( AdaptiveProjectedGuidance, + AdaptiveProjectedMixGuidance, AutoGuidance, + BaseGuidance, ClassifierFreeGuidance, ClassifierFreeZeroStarGuidance, FrequencyDecoupledGuidance, @@ -880,6 +889,8 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: AutoencoderKLAllegro, AutoencoderKLCogVideoX, AutoencoderKLCosmos, + AutoencoderKLHunyuanImage, + AutoencoderKLHunyuanImageRefiner, AutoencoderKLHunyuanVideo, AutoencoderKLLTXVideo, AutoencoderKLMagvit, @@ -912,6 +923,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: HunyuanDiT2DControlNetModel, HunyuanDiT2DModel, HunyuanDiT2DMultiControlNetModel, + HunyuanImageTransformer2DModel, HunyuanVideoFramepackTransformer3DModel, HunyuanVideoTransformer3DModel, I2VGenXLUNet, @@ -1128,6 +1140,8 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: HunyuanDiTControlNetPipeline, HunyuanDiTPAGPipeline, HunyuanDiTPipeline, + HunyuanImagePipeline, + HunyuanImageRefinerPipeline, HunyuanSkyreelsImageToVideoPipeline, HunyuanVideoFramepackPipeline, HunyuanVideoImageToVideoPipeline, diff --git a/src/diffusers/guiders/__init__.py b/src/diffusers/guiders/__init__.py index 23cb7a0a71..6c95b8e0c4 100644 --- a/src/diffusers/guiders/__init__.py +++ b/src/diffusers/guiders/__init__.py @@ -14,28 +14,24 @@ from typing import Union -from ..utils import is_torch_available +from ..utils import is_torch_available, logging + + +logger = logging.get_logger(__name__) +logger.warning( + "Guiders are currently an experimental feature under active development. The API is subject to breaking changes in future releases." +) if is_torch_available(): from .adaptive_projected_guidance import AdaptiveProjectedGuidance + from .adaptive_projected_guidance_mix import AdaptiveProjectedMixGuidance from .auto_guidance import AutoGuidance from .classifier_free_guidance import ClassifierFreeGuidance from .classifier_free_zero_star_guidance import ClassifierFreeZeroStarGuidance from .frequency_decoupled_guidance import FrequencyDecoupledGuidance + from .guider_utils import BaseGuidance from .perturbed_attention_guidance import PerturbedAttentionGuidance from .skip_layer_guidance import SkipLayerGuidance from .smoothed_energy_guidance import SmoothedEnergyGuidance from .tangential_classifier_free_guidance import TangentialClassifierFreeGuidance - - GuiderType = Union[ - AdaptiveProjectedGuidance, - AutoGuidance, - ClassifierFreeGuidance, - ClassifierFreeZeroStarGuidance, - FrequencyDecoupledGuidance, - PerturbedAttentionGuidance, - SkipLayerGuidance, - SmoothedEnergyGuidance, - TangentialClassifierFreeGuidance, - ] diff --git a/src/diffusers/guiders/adaptive_projected_guidance.py b/src/diffusers/guiders/adaptive_projected_guidance.py index 92b1fd5a1c..492d10d2f1 100644 --- a/src/diffusers/guiders/adaptive_projected_guidance.py +++ b/src/diffusers/guiders/adaptive_projected_guidance.py @@ -13,7 +13,7 @@ # limitations under the License. import math -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple import torch @@ -65,8 +65,9 @@ class AdaptiveProjectedGuidance(BaseGuidance): use_original_formulation: bool = False, start: float = 0.0, stop: float = 1.0, + enabled: bool = True, ): - super().__init__(start, stop) + super().__init__(start, stop, enabled) self.guidance_scale = guidance_scale self.adaptive_projected_guidance_momentum = adaptive_projected_guidance_momentum @@ -76,19 +77,14 @@ class AdaptiveProjectedGuidance(BaseGuidance): self.use_original_formulation = use_original_formulation self.momentum_buffer = None - def prepare_inputs( - self, data: "BlockState", input_fields: Optional[Dict[str, Union[str, Tuple[str, str]]]] = None - ) -> List["BlockState"]: - if input_fields is None: - input_fields = self._input_fields - + def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]: if self._step == 0: if self.adaptive_projected_guidance_momentum is not None: self.momentum_buffer = MomentumBuffer(self.adaptive_projected_guidance_momentum) tuple_indices = [0] if self.num_conditions == 1 else [0, 1] data_batches = [] - for i in range(self.num_conditions): - data_batch = self._prepare_batch(input_fields, data, tuple_indices[i], self._input_predictions[i]) + for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions): + data_batch = self._prepare_batch(data, tuple_idx, input_prediction) data_batches.append(data_batch) return data_batches @@ -152,6 +148,44 @@ class MomentumBuffer: new_average = self.momentum * self.running_average self.running_average = update_value + new_average + def __repr__(self) -> str: + """ + Returns a string representation showing momentum, shape, statistics, and a slice of the running_average. + """ + if isinstance(self.running_average, torch.Tensor): + shape = tuple(self.running_average.shape) + + # Calculate statistics + with torch.no_grad(): + stats = { + "mean": self.running_average.mean().item(), + "std": self.running_average.std().item(), + "min": self.running_average.min().item(), + "max": self.running_average.max().item(), + } + + # Get a slice (max 3 elements per dimension) + slice_indices = tuple(slice(None, min(3, dim)) for dim in shape) + sliced_data = self.running_average[slice_indices] + + # Format the slice for display (convert to float32 for numpy compatibility with bfloat16) + slice_str = str(sliced_data.detach().float().cpu().numpy()) + if len(slice_str) > 200: # Truncate if too long + slice_str = slice_str[:200] + "..." + + stats_str = ", ".join([f"{k}={v:.4f}" for k, v in stats.items()]) + + return ( + f"MomentumBuffer(\n" + f" momentum={self.momentum},\n" + f" shape={shape},\n" + f" stats=[{stats_str}],\n" + f" slice={slice_str}\n" + f")" + ) + else: + return f"MomentumBuffer(momentum={self.momentum}, running_average={self.running_average})" + def normalized_guidance( pred_cond: torch.Tensor, diff --git a/src/diffusers/guiders/adaptive_projected_guidance_mix.py b/src/diffusers/guiders/adaptive_projected_guidance_mix.py new file mode 100644 index 0000000000..732741fc92 --- /dev/null +++ b/src/diffusers/guiders/adaptive_projected_guidance_mix.py @@ -0,0 +1,284 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple + +import torch + +from ..configuration_utils import register_to_config +from .guider_utils import BaseGuidance, GuiderOutput, rescale_noise_cfg + + +if TYPE_CHECKING: + from ..modular_pipelines.modular_pipeline import BlockState + + +class AdaptiveProjectedMixGuidance(BaseGuidance): + """ + Adaptive Projected Guidance (APG) https://huggingface.co/papers/2410.02416 combined with Classifier-Free Guidance + (CFG). This guider is used in HunyuanImage2.1 https://github.com/Tencent-Hunyuan/HunyuanImage-2.1 + + Args: + guidance_scale (`float`, defaults to `7.5`): + The scale parameter for classifier-free guidance. Higher values result in stronger conditioning on the text + prompt, while lower values allow for more freedom in generation. Higher values may lead to saturation and + deterioration of image quality. + adaptive_projected_guidance_momentum (`float`, defaults to `None`): + The momentum parameter for the adaptive projected guidance. Disabled if set to `None`. + adaptive_projected_guidance_rescale (`float`, defaults to `15.0`): + The rescale factor applied to the noise predictions for adaptive projected guidance. This is used to + improve image quality and fix + guidance_rescale (`float`, defaults to `0.0`): + The rescale factor applied to the noise predictions for classifier-free guidance. This is used to improve + image quality and fix overexposure. Based on Section 3.4 from [Common Diffusion Noise Schedules and Sample + Steps are Flawed](https://huggingface.co/papers/2305.08891). + use_original_formulation (`bool`, defaults to `False`): + Whether to use the original formulation of classifier-free guidance as proposed in the paper. By default, + we use the diffusers-native implementation that has been in the codebase for a long time. See + [~guiders.classifier_free_guidance.ClassifierFreeGuidance] for more details. + start (`float`, defaults to `0.0`): + The fraction of the total number of denoising steps after which the classifier-free guidance starts. + stop (`float`, defaults to `1.0`): + The fraction of the total number of denoising steps after which the classifier-free guidance stops. + adaptive_projected_guidance_start_step (`int`, defaults to `5`): + The step at which the adaptive projected guidance starts (before this step, classifier-free guidance is + used, and momentum buffer is updated). + enabled (`bool`, defaults to `True`): + Whether this guidance is enabled. + """ + + _input_predictions = ["pred_cond", "pred_uncond"] + + @register_to_config + def __init__( + self, + guidance_scale: float = 3.5, + guidance_rescale: float = 0.0, + adaptive_projected_guidance_scale: float = 10.0, + adaptive_projected_guidance_momentum: float = -0.5, + adaptive_projected_guidance_rescale: float = 10.0, + eta: float = 0.0, + use_original_formulation: bool = False, + start: float = 0.0, + stop: float = 1.0, + adaptive_projected_guidance_start_step: int = 5, + enabled: bool = True, + ): + super().__init__(start, stop, enabled) + + self.guidance_scale = guidance_scale + self.guidance_rescale = guidance_rescale + self.adaptive_projected_guidance_scale = adaptive_projected_guidance_scale + self.adaptive_projected_guidance_momentum = adaptive_projected_guidance_momentum + self.adaptive_projected_guidance_rescale = adaptive_projected_guidance_rescale + self.eta = eta + self.adaptive_projected_guidance_start_step = adaptive_projected_guidance_start_step + self.use_original_formulation = use_original_formulation + self.momentum_buffer = None + + def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]: + if self._step == 0: + if self.adaptive_projected_guidance_momentum is not None: + self.momentum_buffer = MomentumBuffer(self.adaptive_projected_guidance_momentum) + tuple_indices = [0] if self.num_conditions == 1 else [0, 1] + data_batches = [] + for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions): + data_batch = self._prepare_batch(data, tuple_idx, input_prediction) + data_batches.append(data_batch) + return data_batches + + def forward(self, pred_cond: torch.Tensor, pred_uncond: Optional[torch.Tensor] = None) -> GuiderOutput: + pred = None + + # no guidance + if not self._is_cfg_enabled(): + pred = pred_cond + + # CFG + update momentum buffer + elif not self._is_apg_enabled(): + if self.momentum_buffer is not None: + update_momentum_buffer(pred_cond, pred_uncond, self.momentum_buffer) + # CFG + update momentum buffer + shift = pred_cond - pred_uncond + pred = pred_cond if self.use_original_formulation else pred_uncond + pred = pred + self.guidance_scale * shift + + # APG + elif self._is_apg_enabled(): + pred = normalized_guidance( + pred_cond, + pred_uncond, + self.adaptive_projected_guidance_scale, + self.momentum_buffer, + self.eta, + self.adaptive_projected_guidance_rescale, + self.use_original_formulation, + ) + + if self.guidance_rescale > 0.0: + pred = rescale_noise_cfg(pred, pred_cond, self.guidance_rescale) + + return GuiderOutput(pred=pred, pred_cond=pred_cond, pred_uncond=pred_uncond) + + @property + def is_conditional(self) -> bool: + return self._count_prepared == 1 + + @property + def num_conditions(self) -> int: + num_conditions = 1 + if self._is_apg_enabled() or self._is_cfg_enabled(): + num_conditions += 1 + return num_conditions + + # Copied from diffusers.guiders.classifier_free_guidance.ClassifierFreeGuidance._is_cfg_enabled + def _is_cfg_enabled(self) -> bool: + if not self._enabled: + return False + + is_within_range = True + if self._num_inference_steps is not None: + skip_start_step = int(self._start * self._num_inference_steps) + skip_stop_step = int(self._stop * self._num_inference_steps) + is_within_range = skip_start_step <= self._step < skip_stop_step + + is_close = False + if self.use_original_formulation: + is_close = math.isclose(self.guidance_scale, 0.0) + else: + is_close = math.isclose(self.guidance_scale, 1.0) + + return is_within_range and not is_close + + def _is_apg_enabled(self) -> bool: + if not self._enabled: + return False + + if not self._is_cfg_enabled(): + return False + + is_within_range = False + if self._step is not None: + is_within_range = self._step > self.adaptive_projected_guidance_start_step + + is_close = False + if self.use_original_formulation: + is_close = math.isclose(self.adaptive_projected_guidance_scale, 0.0) + else: + is_close = math.isclose(self.adaptive_projected_guidance_scale, 1.0) + + return is_within_range and not is_close + + def get_state(self): + state = super().get_state() + state["momentum_buffer"] = self.momentum_buffer + state["is_apg_enabled"] = self._is_apg_enabled() + state["is_cfg_enabled"] = self._is_cfg_enabled() + return state + + +# Copied from diffusers.guiders.adaptive_projected_guidance.MomentumBuffer +class MomentumBuffer: + def __init__(self, momentum: float): + self.momentum = momentum + self.running_average = 0 + + def update(self, update_value: torch.Tensor): + new_average = self.momentum * self.running_average + self.running_average = update_value + new_average + + def __repr__(self) -> str: + """ + Returns a string representation showing momentum, shape, statistics, and a slice of the running_average. + """ + if isinstance(self.running_average, torch.Tensor): + shape = tuple(self.running_average.shape) + + # Calculate statistics + with torch.no_grad(): + stats = { + "mean": self.running_average.mean().item(), + "std": self.running_average.std().item(), + "min": self.running_average.min().item(), + "max": self.running_average.max().item(), + } + + # Get a slice (max 3 elements per dimension) + slice_indices = tuple(slice(None, min(3, dim)) for dim in shape) + sliced_data = self.running_average[slice_indices] + + # Format the slice for display (convert to float32 for numpy compatibility with bfloat16) + slice_str = str(sliced_data.detach().float().cpu().numpy()) + if len(slice_str) > 200: # Truncate if too long + slice_str = slice_str[:200] + "..." + + stats_str = ", ".join([f"{k}={v:.4f}" for k, v in stats.items()]) + + return ( + f"MomentumBuffer(\n" + f" momentum={self.momentum},\n" + f" shape={shape},\n" + f" stats=[{stats_str}],\n" + f" slice={slice_str}\n" + f")" + ) + else: + return f"MomentumBuffer(momentum={self.momentum}, running_average={self.running_average})" + + +def update_momentum_buffer( + pred_cond: torch.Tensor, + pred_uncond: torch.Tensor, + momentum_buffer: Optional[MomentumBuffer] = None, +): + diff = pred_cond - pred_uncond + if momentum_buffer is not None: + momentum_buffer.update(diff) + + +def normalized_guidance( + pred_cond: torch.Tensor, + pred_uncond: torch.Tensor, + guidance_scale: float, + momentum_buffer: Optional[MomentumBuffer] = None, + eta: float = 1.0, + norm_threshold: float = 0.0, + use_original_formulation: bool = False, +): + if momentum_buffer is not None: + update_momentum_buffer(pred_cond, pred_uncond, momentum_buffer) + diff = momentum_buffer.running_average + else: + diff = pred_cond - pred_uncond + + dim = [-i for i in range(1, len(diff.shape))] + + if norm_threshold > 0: + ones = torch.ones_like(diff) + diff_norm = diff.norm(p=2, dim=dim, keepdim=True) + scale_factor = torch.minimum(ones, norm_threshold / diff_norm) + diff = diff * scale_factor + + v0, v1 = diff.double(), pred_cond.double() + v1 = torch.nn.functional.normalize(v1, dim=dim) + v0_parallel = (v0 * v1).sum(dim=dim, keepdim=True) * v1 + v0_orthogonal = v0 - v0_parallel + diff_parallel, diff_orthogonal = v0_parallel.type_as(diff), v0_orthogonal.type_as(diff) + normalized_update = diff_orthogonal + eta * diff_parallel + + pred = pred_cond if use_original_formulation else pred_uncond + pred = pred + guidance_scale * normalized_update + + return pred diff --git a/src/diffusers/guiders/auto_guidance.py b/src/diffusers/guiders/auto_guidance.py index 5271a530ea..4374f45aff 100644 --- a/src/diffusers/guiders/auto_guidance.py +++ b/src/diffusers/guiders/auto_guidance.py @@ -72,8 +72,9 @@ class AutoGuidance(BaseGuidance): use_original_formulation: bool = False, start: float = 0.0, stop: float = 1.0, + enabled: bool = True, ): - super().__init__(start, stop) + super().__init__(start, stop, enabled) self.guidance_scale = guidance_scale self.auto_guidance_layers = auto_guidance_layers @@ -132,16 +133,11 @@ class AutoGuidance(BaseGuidance): registry = HookRegistry.check_if_exists_or_initialize(denoiser) registry.remove_hook(name, recurse=True) - def prepare_inputs( - self, data: "BlockState", input_fields: Optional[Dict[str, Union[str, Tuple[str, str]]]] = None - ) -> List["BlockState"]: - if input_fields is None: - input_fields = self._input_fields - + def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]: tuple_indices = [0] if self.num_conditions == 1 else [0, 1] data_batches = [] - for i in range(self.num_conditions): - data_batch = self._prepare_batch(input_fields, data, tuple_indices[i], self._input_predictions[i]) + for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions): + data_batch = self._prepare_batch(data, tuple_idx, input_prediction) data_batches.append(data_batch) return data_batches diff --git a/src/diffusers/guiders/classifier_free_guidance.py b/src/diffusers/guiders/classifier_free_guidance.py index 050590336f..d475b30226 100644 --- a/src/diffusers/guiders/classifier_free_guidance.py +++ b/src/diffusers/guiders/classifier_free_guidance.py @@ -13,7 +13,7 @@ # limitations under the License. import math -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple import torch @@ -27,43 +27,50 @@ if TYPE_CHECKING: class ClassifierFreeGuidance(BaseGuidance): """ - Classifier-free guidance (CFG): https://huggingface.co/papers/2207.12598 + Implements Classifier-Free Guidance (CFG) for diffusion models. - CFG is a technique used to improve generation quality and condition-following in diffusion models. It works by - jointly training a model on both conditional and unconditional data, and using a weighted sum of the two during - inference. This allows the model to tradeoff between generation quality and sample diversity. The original paper - proposes scaling and shifting the conditional distribution based on the difference between conditional and - unconditional predictions. [x_pred = x_cond + scale * (x_cond - x_uncond)] + Reference: https://huggingface.co/papers/2207.12598 - Diffusers implemented the scaling and shifting on the unconditional prediction instead based on the [Imagen - paper](https://huggingface.co/papers/2205.11487), which is equivalent to what the original paper proposed in - theory. [x_pred = x_uncond + scale * (x_cond - x_uncond)] + CFG improves generation quality and prompt adherence by jointly training models on both conditional and + unconditional data, then combining predictions during inference. This allows trading off between quality (high + guidance) and diversity (low guidance). - The intution behind the original formulation can be thought of as moving the conditional distribution estimates - further away from the unconditional distribution estimates, while the diffusers-native implementation can be - thought of as moving the unconditional distribution towards the conditional distribution estimates to get rid of - the unconditional predictions (usually negative features like "bad quality, bad anotomy, watermarks", etc.) + **Two CFG Formulations:** - The `use_original_formulation` argument can be set to `True` to use the original CFG formulation mentioned in the - paper. By default, we use the diffusers-native implementation that has been in the codebase for a long time. + 1. **Original formulation** (from paper): + ``` + x_pred = x_cond + guidance_scale * (x_cond - x_uncond) + ``` + Moves conditional predictions further from unconditional ones. + + 2. **Diffusers-native formulation** (default, from Imagen paper): + ``` + x_pred = x_uncond + guidance_scale * (x_cond - x_uncond) + ``` + Moves unconditional predictions toward conditional ones, effectively suppressing negative features (e.g., "bad + quality", "watermarks"). Equivalent in theory but more intuitive. + + Use `use_original_formulation=True` to switch to the original formulation. Args: guidance_scale (`float`, defaults to `7.5`): - The scale parameter for classifier-free guidance. Higher values result in stronger conditioning on the text - prompt, while lower values allow for more freedom in generation. Higher values may lead to saturation and - deterioration of image quality. + CFG scale applied by this guider during post-processing. Higher values = stronger prompt conditioning but + may reduce quality. Typical range: 1.0-20.0. guidance_rescale (`float`, defaults to `0.0`): - The rescale factor applied to the noise predictions. This is used to improve image quality and fix - overexposure. Based on Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are - Flawed](https://huggingface.co/papers/2305.08891). + Rescaling factor to prevent overexposure from high guidance scales. Based on [Common Diffusion Noise + Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891). Range: 0.0 (no rescaling) + to 1.0 (full rescaling). use_original_formulation (`bool`, defaults to `False`): - Whether to use the original formulation of classifier-free guidance as proposed in the paper. By default, - we use the diffusers-native implementation that has been in the codebase for a long time. See - [~guiders.classifier_free_guidance.ClassifierFreeGuidance] for more details. + If `True`, uses the original CFG formulation from the paper. If `False` (default), uses the + diffusers-native formulation from the Imagen paper. start (`float`, defaults to `0.0`): - The fraction of the total number of denoising steps after which guidance starts. + Fraction of denoising steps (0.0-1.0) after which CFG starts. Use > 0.0 to disable CFG in early denoising + steps. stop (`float`, defaults to `1.0`): - The fraction of the total number of denoising steps after which guidance stops. + Fraction of denoising steps (0.0-1.0) after which CFG stops. Use < 1.0 to disable CFG in late denoising + steps. + enabled (`bool`, defaults to `True`): + Whether CFG is enabled. Set to `False` to disable CFG entirely (uses only conditional predictions). """ _input_predictions = ["pred_cond", "pred_uncond"] @@ -76,23 +83,19 @@ class ClassifierFreeGuidance(BaseGuidance): use_original_formulation: bool = False, start: float = 0.0, stop: float = 1.0, + enabled: bool = True, ): - super().__init__(start, stop) + super().__init__(start, stop, enabled) self.guidance_scale = guidance_scale self.guidance_rescale = guidance_rescale self.use_original_formulation = use_original_formulation - def prepare_inputs( - self, data: "BlockState", input_fields: Optional[Dict[str, Union[str, Tuple[str, str]]]] = None - ) -> List["BlockState"]: - if input_fields is None: - input_fields = self._input_fields - + def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]: tuple_indices = [0] if self.num_conditions == 1 else [0, 1] data_batches = [] - for i in range(self.num_conditions): - data_batch = self._prepare_batch(input_fields, data, tuple_indices[i], self._input_predictions[i]) + for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions): + data_batch = self._prepare_batch(data, tuple_idx, input_prediction) data_batches.append(data_batch) return data_batches diff --git a/src/diffusers/guiders/classifier_free_zero_star_guidance.py b/src/diffusers/guiders/classifier_free_zero_star_guidance.py index b64e356331..1ea6bbb1c8 100644 --- a/src/diffusers/guiders/classifier_free_zero_star_guidance.py +++ b/src/diffusers/guiders/classifier_free_zero_star_guidance.py @@ -13,7 +13,7 @@ # limitations under the License. import math -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple import torch @@ -68,31 +68,31 @@ class ClassifierFreeZeroStarGuidance(BaseGuidance): use_original_formulation: bool = False, start: float = 0.0, stop: float = 1.0, + enabled: bool = True, ): - super().__init__(start, stop) + super().__init__(start, stop, enabled) self.guidance_scale = guidance_scale self.zero_init_steps = zero_init_steps self.guidance_rescale = guidance_rescale self.use_original_formulation = use_original_formulation - def prepare_inputs( - self, data: "BlockState", input_fields: Optional[Dict[str, Union[str, Tuple[str, str]]]] = None - ) -> List["BlockState"]: - if input_fields is None: - input_fields = self._input_fields - + def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]: tuple_indices = [0] if self.num_conditions == 1 else [0, 1] data_batches = [] - for i in range(self.num_conditions): - data_batch = self._prepare_batch(input_fields, data, tuple_indices[i], self._input_predictions[i]) + for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions): + data_batch = self._prepare_batch(data, tuple_idx, input_prediction) data_batches.append(data_batch) return data_batches def forward(self, pred_cond: torch.Tensor, pred_uncond: Optional[torch.Tensor] = None) -> GuiderOutput: pred = None - if self._step < self.zero_init_steps: + # YiYi Notes: add default behavior for self._enabled == False + if not self._enabled: + pred = pred_cond + + elif self._step < self.zero_init_steps: pred = torch.zeros_like(pred_cond) elif not self._is_cfg_enabled(): pred = pred_cond diff --git a/src/diffusers/guiders/frequency_decoupled_guidance.py b/src/diffusers/guiders/frequency_decoupled_guidance.py index 93822a180e..cd542a43a4 100644 --- a/src/diffusers/guiders/frequency_decoupled_guidance.py +++ b/src/diffusers/guiders/frequency_decoupled_guidance.py @@ -149,6 +149,7 @@ class FrequencyDecoupledGuidance(BaseGuidance): stop: Union[float, List[float], Tuple[float]] = 1.0, guidance_rescale_space: str = "data", upcast_to_double: bool = True, + enabled: bool = True, ): if not _CAN_USE_KORNIA: raise ImportError( @@ -160,7 +161,7 @@ class FrequencyDecoupledGuidance(BaseGuidance): # Set start to earliest start for any freq component and stop to latest stop for any freq component min_start = start if isinstance(start, float) else min(start) max_stop = stop if isinstance(stop, float) else max(stop) - super().__init__(min_start, max_stop) + super().__init__(min_start, max_stop, enabled) self.guidance_scales = guidance_scales self.levels = len(guidance_scales) @@ -217,16 +218,11 @@ class FrequencyDecoupledGuidance(BaseGuidance): f"({len(self.guidance_scales)})" ) - def prepare_inputs( - self, data: "BlockState", input_fields: Optional[Dict[str, Union[str, Tuple[str, str]]]] = None - ) -> List["BlockState"]: - if input_fields is None: - input_fields = self._input_fields - + def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]: tuple_indices = [0] if self.num_conditions == 1 else [0, 1] data_batches = [] - for i in range(self.num_conditions): - data_batch = self._prepare_batch(input_fields, data, tuple_indices[i], self._input_predictions[i]) + for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions): + data_batch = self._prepare_batch(data, tuple_idx, input_prediction) data_batches.append(data_batch) return data_batches diff --git a/src/diffusers/guiders/guider_utils.py b/src/diffusers/guiders/guider_utils.py index 7524b5a3ea..b7afd8648b 100644 --- a/src/diffusers/guiders/guider_utils.py +++ b/src/diffusers/guiders/guider_utils.py @@ -40,7 +40,7 @@ class BaseGuidance(ConfigMixin, PushToHubMixin): _input_predictions = None _identifier_key = "__guidance_identifier__" - def __init__(self, start: float = 0.0, stop: float = 1.0): + def __init__(self, start: float = 0.0, stop: float = 1.0, enabled: bool = True): self._start = start self._stop = stop self._step: int = None @@ -48,7 +48,7 @@ class BaseGuidance(ConfigMixin, PushToHubMixin): self._timestep: torch.LongTensor = None self._count_prepared = 0 self._input_fields: Dict[str, Union[str, Tuple[str, str]]] = None - self._enabled = True + self._enabled = enabled if not (0.0 <= start < 1.0): raise ValueError(f"Expected `start` to be between 0.0 and 1.0, but got {start}.") @@ -60,6 +60,31 @@ class BaseGuidance(ConfigMixin, PushToHubMixin): "`_input_predictions` must be a list of required prediction names for the guidance technique." ) + def new(self, **kwargs): + """ + Creates a copy of this guider instance, optionally with modified configuration parameters. + + Args: + **kwargs: Configuration parameters to override in the new instance. If no kwargs are provided, + returns an exact copy with the same configuration. + + Returns: + A new guider instance with the same (or updated) configuration. + + Example: + ```python + # Create a CFG guider + guider = ClassifierFreeGuidance(guidance_scale=3.5) + + # Create an exact copy + same_guider = guider.new() + + # Create a copy with different start step, keeping other config the same + new_guider = guider.new(guidance_scale=5) + ``` + """ + return self.__class__.from_config(self.config, **kwargs) + def disable(self): self._enabled = False @@ -72,42 +97,52 @@ class BaseGuidance(ConfigMixin, PushToHubMixin): self._timestep = timestep self._count_prepared = 0 - def set_input_fields(self, **kwargs: Dict[str, Union[str, Tuple[str, str]]]) -> None: + def get_state(self) -> Dict[str, Any]: """ - Set the input fields for the guidance technique. The input fields are used to specify the names of the returned - attributes containing the prepared data after `prepare_inputs` is called. The prepared data is obtained from - the values of the provided keyword arguments to this method. - - Args: - **kwargs (`Dict[str, Union[str, Tuple[str, str]]]`): - A dictionary where the keys are the names of the fields that will be used to store the data once it is - prepared with `prepare_inputs`. The values can be either a string or a tuple of length 2, which is used - to look up the required data provided for preparation. - - If a string is provided, it will be used as the conditional data (or unconditional if used with a - guidance method that requires it). If a tuple of length 2 is provided, the first element must be the - conditional data identifier and the second element must be the unconditional data identifier or None. - - Example: - ``` - data = {"prompt_embeds": , "negative_prompt_embeds": , "latents": } - - BaseGuidance.set_input_fields( - latents="latents", - prompt_embeds=("prompt_embeds", "negative_prompt_embeds"), - ) - ``` + Returns the current state of the guidance technique as a dictionary. The state variables will be included in + the __repr__ method. Returns: + `Dict[str, Any]`: A dictionary containing the current state variables including: + - step: Current inference step + - num_inference_steps: Total number of inference steps + - timestep: Current timestep tensor + - count_prepared: Number of times prepare_models has been called + - enabled: Whether the guidance is enabled + - num_conditions: Number of conditions """ - for key, value in kwargs.items(): - is_string = isinstance(value, str) - is_tuple_of_str_with_len_2 = ( - isinstance(value, tuple) and len(value) == 2 and all(isinstance(v, str) for v in value) - ) - if not (is_string or is_tuple_of_str_with_len_2): - raise ValueError( - f"Expected `set_input_fields` to be called with a string or a tuple of string with length 2, but got {type(value)} for key {key}." - ) - self._input_fields = kwargs + state = { + "step": self._step, + "num_inference_steps": self._num_inference_steps, + "timestep": self._timestep, + "count_prepared": self._count_prepared, + "enabled": self._enabled, + "num_conditions": self.num_conditions, + } + return state + + def __repr__(self) -> str: + """ + Returns a string representation of the guidance object including both config and current state. + """ + # Get ConfigMixin's __repr__ + str_repr = super().__repr__() + + # Get current state + state = self.get_state() + + # Format each state variable on its own line with indentation + state_lines = [] + for k, v in state.items(): + # Convert value to string and handle multi-line values + v_str = str(v) + if "\n" in v_str: + # For multi-line values (like MomentumBuffer), indent subsequent lines + v_lines = v_str.split("\n") + v_str = v_lines[0] + "\n" + "\n".join([" " + line for line in v_lines[1:]]) + state_lines.append(f" {k}: {v_str}") + + state_str = "\n".join(state_lines) + + return f"{str_repr}\nState:\n{state_str}" def prepare_models(self, denoiser: torch.nn.Module) -> None: """ @@ -155,8 +190,7 @@ class BaseGuidance(ConfigMixin, PushToHubMixin): @classmethod def _prepare_batch( cls, - input_fields: Dict[str, Union[str, Tuple[str, str]]], - data: "BlockState", + data: Dict[str, Tuple[torch.Tensor, torch.Tensor]], tuple_index: int, identifier: str, ) -> "BlockState": @@ -182,21 +216,16 @@ class BaseGuidance(ConfigMixin, PushToHubMixin): """ from ..modular_pipelines.modular_pipeline import BlockState - if input_fields is None: - raise ValueError( - "Input fields cannot be None. Please pass `input_fields` to `prepare_inputs` or call `set_input_fields` before preparing inputs." - ) data_batch = {} - for key, value in input_fields.items(): + for key, value in data.items(): try: - if isinstance(value, str): - data_batch[key] = getattr(data, value) + if isinstance(value, torch.Tensor): + data_batch[key] = value elif isinstance(value, tuple): - data_batch[key] = getattr(data, value[tuple_index]) + data_batch[key] = value[tuple_index] else: - # We've already checked that value is a string or a tuple of strings with length 2 - pass - except AttributeError: + raise ValueError(f"Invalid value type: {type(value)}") + except ValueError: logger.debug(f"`data` does not have attribute(s) {value}, skipping.") data_batch[cls._identifier_key] = identifier return BlockState(**data_batch) diff --git a/src/diffusers/guiders/perturbed_attention_guidance.py b/src/diffusers/guiders/perturbed_attention_guidance.py index e294e8d0db..29341736e8 100644 --- a/src/diffusers/guiders/perturbed_attention_guidance.py +++ b/src/diffusers/guiders/perturbed_attention_guidance.py @@ -98,8 +98,9 @@ class PerturbedAttentionGuidance(BaseGuidance): use_original_formulation: bool = False, start: float = 0.0, stop: float = 1.0, + enabled: bool = True, ): - super().__init__(start, stop) + super().__init__(start, stop, enabled) self.guidance_scale = guidance_scale self.skip_layer_guidance_scale = perturbed_guidance_scale @@ -168,12 +169,7 @@ class PerturbedAttentionGuidance(BaseGuidance): registry.remove_hook(hook_name, recurse=True) # Copied from diffusers.guiders.skip_layer_guidance.SkipLayerGuidance.prepare_inputs - def prepare_inputs( - self, data: "BlockState", input_fields: Optional[Dict[str, Union[str, Tuple[str, str]]]] = None - ) -> List["BlockState"]: - if input_fields is None: - input_fields = self._input_fields - + def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]: if self.num_conditions == 1: tuple_indices = [0] input_predictions = ["pred_cond"] @@ -186,8 +182,8 @@ class PerturbedAttentionGuidance(BaseGuidance): tuple_indices = [0, 1, 0] input_predictions = ["pred_cond", "pred_uncond", "pred_cond_skip"] data_batches = [] - for i in range(self.num_conditions): - data_batch = self._prepare_batch(input_fields, data, tuple_indices[i], input_predictions[i]) + for tuple_idx, input_prediction in zip(tuple_indices, input_predictions): + data_batch = self._prepare_batch(data, tuple_idx, input_prediction) data_batches.append(data_batch) return data_batches diff --git a/src/diffusers/guiders/skip_layer_guidance.py b/src/diffusers/guiders/skip_layer_guidance.py index 3530df8b0a..fa5b93b680 100644 --- a/src/diffusers/guiders/skip_layer_guidance.py +++ b/src/diffusers/guiders/skip_layer_guidance.py @@ -100,8 +100,9 @@ class SkipLayerGuidance(BaseGuidance): use_original_formulation: bool = False, start: float = 0.0, stop: float = 1.0, + enabled: bool = True, ): - super().__init__(start, stop) + super().__init__(start, stop, enabled) self.guidance_scale = guidance_scale self.skip_layer_guidance_scale = skip_layer_guidance_scale @@ -164,12 +165,7 @@ class SkipLayerGuidance(BaseGuidance): for hook_name in self._skip_layer_hook_names: registry.remove_hook(hook_name, recurse=True) - def prepare_inputs( - self, data: "BlockState", input_fields: Optional[Dict[str, Union[str, Tuple[str, str]]]] = None - ) -> List["BlockState"]: - if input_fields is None: - input_fields = self._input_fields - + def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]: if self.num_conditions == 1: tuple_indices = [0] input_predictions = ["pred_cond"] @@ -182,8 +178,8 @@ class SkipLayerGuidance(BaseGuidance): tuple_indices = [0, 1, 0] input_predictions = ["pred_cond", "pred_uncond", "pred_cond_skip"] data_batches = [] - for i in range(self.num_conditions): - data_batch = self._prepare_batch(input_fields, data, tuple_indices[i], input_predictions[i]) + for tuple_idx, input_prediction in zip(tuple_indices, input_predictions): + data_batch = self._prepare_batch(data, tuple_idx, input_prediction) data_batches.append(data_batch) return data_batches diff --git a/src/diffusers/guiders/smoothed_energy_guidance.py b/src/diffusers/guiders/smoothed_energy_guidance.py index 767d20b62f..7446b33f12 100644 --- a/src/diffusers/guiders/smoothed_energy_guidance.py +++ b/src/diffusers/guiders/smoothed_energy_guidance.py @@ -92,8 +92,9 @@ class SmoothedEnergyGuidance(BaseGuidance): use_original_formulation: bool = False, start: float = 0.0, stop: float = 1.0, + enabled: bool = True, ): - super().__init__(start, stop) + super().__init__(start, stop, enabled) self.guidance_scale = guidance_scale self.seg_guidance_scale = seg_guidance_scale @@ -153,12 +154,7 @@ class SmoothedEnergyGuidance(BaseGuidance): for hook_name in self._seg_layer_hook_names: registry.remove_hook(hook_name, recurse=True) - def prepare_inputs( - self, data: "BlockState", input_fields: Optional[Dict[str, Union[str, Tuple[str, str]]]] = None - ) -> List["BlockState"]: - if input_fields is None: - input_fields = self._input_fields - + def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]: if self.num_conditions == 1: tuple_indices = [0] input_predictions = ["pred_cond"] @@ -171,8 +167,8 @@ class SmoothedEnergyGuidance(BaseGuidance): tuple_indices = [0, 1, 0] input_predictions = ["pred_cond", "pred_uncond", "pred_cond_seg"] data_batches = [] - for i in range(self.num_conditions): - data_batch = self._prepare_batch(input_fields, data, tuple_indices[i], input_predictions[i]) + for tuple_idx, input_prediction in zip(tuple_indices, input_predictions): + data_batch = self._prepare_batch(data, tuple_idx, input_prediction) data_batches.append(data_batch) return data_batches diff --git a/src/diffusers/guiders/tangential_classifier_free_guidance.py b/src/diffusers/guiders/tangential_classifier_free_guidance.py index df1e69fe71..cfa3c4a616 100644 --- a/src/diffusers/guiders/tangential_classifier_free_guidance.py +++ b/src/diffusers/guiders/tangential_classifier_free_guidance.py @@ -13,7 +13,7 @@ # limitations under the License. import math -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple import torch @@ -58,23 +58,19 @@ class TangentialClassifierFreeGuidance(BaseGuidance): use_original_formulation: bool = False, start: float = 0.0, stop: float = 1.0, + enabled: bool = True, ): - super().__init__(start, stop) + super().__init__(start, stop, enabled) self.guidance_scale = guidance_scale self.guidance_rescale = guidance_rescale self.use_original_formulation = use_original_formulation - def prepare_inputs( - self, data: "BlockState", input_fields: Optional[Dict[str, Union[str, Tuple[str, str]]]] = None - ) -> List["BlockState"]: - if input_fields is None: - input_fields = self._input_fields - + def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]: tuple_indices = [0] if self.num_conditions == 1 else [0, 1] data_batches = [] - for i in range(self.num_conditions): - data_batch = self._prepare_batch(input_fields, data, tuple_indices[i], self._input_predictions[i]) + for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions): + data_batch = self._prepare_batch(data, tuple_idx, input_prediction) data_batches.append(data_batch) return data_batches diff --git a/src/diffusers/hooks/_helpers.py b/src/diffusers/hooks/_helpers.py index f6e5bdd52d..790199f3c9 100644 --- a/src/diffusers/hooks/_helpers.py +++ b/src/diffusers/hooks/_helpers.py @@ -108,6 +108,7 @@ def _register_attention_processors_metadata(): from ..models.attention_processor import AttnProcessor2_0 from ..models.transformers.transformer_cogview4 import CogView4AttnProcessor from ..models.transformers.transformer_flux import FluxAttnProcessor + from ..models.transformers.transformer_hunyuanimage import HunyuanImageAttnProcessor from ..models.transformers.transformer_qwenimage import QwenDoubleStreamAttnProcessor2_0 from ..models.transformers.transformer_wan import WanAttnProcessor2_0 @@ -149,6 +150,14 @@ def _register_attention_processors_metadata(): ), ) + # HunyuanImageAttnProcessor + AttentionProcessorRegistry.register( + model_class=HunyuanImageAttnProcessor, + metadata=AttentionProcessorMetadata( + skip_processor_output_fn=_skip_proc_output_fn_Attention_HunyuanImageAttnProcessor, + ), + ) + def _register_transformer_blocks_metadata(): from ..models.attention import BasicTransformerBlock @@ -162,6 +171,10 @@ def _register_transformer_blocks_metadata(): HunyuanVideoTokenReplaceTransformerBlock, HunyuanVideoTransformerBlock, ) + from ..models.transformers.transformer_hunyuanimage import ( + HunyuanImageSingleTransformerBlock, + HunyuanImageTransformerBlock, + ) from ..models.transformers.transformer_ltx import LTXVideoTransformerBlock from ..models.transformers.transformer_mochi import MochiTransformerBlock from ..models.transformers.transformer_qwenimage import QwenImageTransformerBlock @@ -283,6 +296,22 @@ def _register_transformer_blocks_metadata(): ), ) + # HunyuanImage2.1 + TransformerBlockRegistry.register( + model_class=HunyuanImageTransformerBlock, + metadata=TransformerBlockMetadata( + return_hidden_states_index=0, + return_encoder_hidden_states_index=1, + ), + ) + TransformerBlockRegistry.register( + model_class=HunyuanImageSingleTransformerBlock, + metadata=TransformerBlockMetadata( + return_hidden_states_index=0, + return_encoder_hidden_states_index=1, + ), + ) + # fmt: off def _skip_attention___ret___hidden_states(self, *args, **kwargs): @@ -308,4 +337,5 @@ _skip_proc_output_fn_Attention_WanAttnProcessor2_0 = _skip_attention___ret___hid # not sure what this is yet. _skip_proc_output_fn_Attention_FluxAttnProcessor = _skip_attention___ret___hidden_states _skip_proc_output_fn_Attention_QwenDoubleStreamAttnProcessor2_0 = _skip_attention___ret___hidden_states +_skip_proc_output_fn_Attention_HunyuanImageAttnProcessor = _skip_attention___ret___hidden_states # fmt: on diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index 532db76a09..241ccf7b78 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -36,6 +36,8 @@ if is_torch_available(): _import_structure["autoencoders.autoencoder_kl_cogvideox"] = ["AutoencoderKLCogVideoX"] _import_structure["autoencoders.autoencoder_kl_cosmos"] = ["AutoencoderKLCosmos"] _import_structure["autoencoders.autoencoder_kl_hunyuan_video"] = ["AutoencoderKLHunyuanVideo"] + _import_structure["autoencoders.autoencoder_kl_hunyuanimage"] = ["AutoencoderKLHunyuanImage"] + _import_structure["autoencoders.autoencoder_kl_hunyuanimage_refiner"] = ["AutoencoderKLHunyuanImageRefiner"] _import_structure["autoencoders.autoencoder_kl_ltx"] = ["AutoencoderKLLTXVideo"] _import_structure["autoencoders.autoencoder_kl_magvit"] = ["AutoencoderKLMagvit"] _import_structure["autoencoders.autoencoder_kl_mochi"] = ["AutoencoderKLMochi"] @@ -91,6 +93,7 @@ if is_torch_available(): _import_structure["transformers.transformer_hidream_image"] = ["HiDreamImageTransformer2DModel"] _import_structure["transformers.transformer_hunyuan_video"] = ["HunyuanVideoTransformer3DModel"] _import_structure["transformers.transformer_hunyuan_video_framepack"] = ["HunyuanVideoFramepackTransformer3DModel"] + _import_structure["transformers.transformer_hunyuanimage"] = ["HunyuanImageTransformer2DModel"] _import_structure["transformers.transformer_kandinsky"] = ["Kandinsky5Transformer3DModel"] _import_structure["transformers.transformer_ltx"] = ["LTXVideoTransformer3DModel"] _import_structure["transformers.transformer_lumina2"] = ["Lumina2Transformer2DModel"] @@ -133,6 +136,8 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: AutoencoderKLAllegro, AutoencoderKLCogVideoX, AutoencoderKLCosmos, + AutoencoderKLHunyuanImage, + AutoencoderKLHunyuanImageRefiner, AutoencoderKLHunyuanVideo, AutoencoderKLLTXVideo, AutoencoderKLMagvit, @@ -182,6 +187,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: FluxTransformer2DModel, HiDreamImageTransformer2DModel, HunyuanDiT2DModel, + HunyuanImageTransformer2DModel, HunyuanVideoFramepackTransformer3DModel, HunyuanVideoTransformer3DModel, Kandinsky5Transformer3DModel, diff --git a/src/diffusers/models/autoencoders/__init__.py b/src/diffusers/models/autoencoders/__init__.py index c008a45298..edfaabb070 100644 --- a/src/diffusers/models/autoencoders/__init__.py +++ b/src/diffusers/models/autoencoders/__init__.py @@ -5,6 +5,8 @@ from .autoencoder_kl_allegro import AutoencoderKLAllegro from .autoencoder_kl_cogvideox import AutoencoderKLCogVideoX from .autoencoder_kl_cosmos import AutoencoderKLCosmos from .autoencoder_kl_hunyuan_video import AutoencoderKLHunyuanVideo +from .autoencoder_kl_hunyuanimage import AutoencoderKLHunyuanImage +from .autoencoder_kl_hunyuanimage_refiner import AutoencoderKLHunyuanImageRefiner from .autoencoder_kl_ltx import AutoencoderKLLTXVideo from .autoencoder_kl_magvit import AutoencoderKLMagvit from .autoencoder_kl_mochi import AutoencoderKLMochi diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage.py b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage.py new file mode 100644 index 0000000000..616d0d4158 --- /dev/null +++ b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage.py @@ -0,0 +1,709 @@ +# Copyright 2025 The Hunyuan Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Optional, Tuple, Union + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin +from ...utils import logging +from ...utils.accelerate_utils import apply_forward_hook +from ..activations import get_activation +from ..modeling_outputs import AutoencoderKLOutput +from ..modeling_utils import ModelMixin +from .vae import DecoderOutput, DiagonalGaussianDistribution + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class HunyuanImageResnetBlock(nn.Module): + r""" + Residual block with two convolutions and optional channel change. + + Args: + in_channels (int): Number of input channels. + out_channels (int): Number of output channels. + non_linearity (str, optional): Type of non-linearity to use. Default is "silu". + """ + + def __init__(self, in_channels: int, out_channels: int, non_linearity: str = "silu") -> None: + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.nonlinearity = get_activation(non_linearity) + + # layers + self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True) + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) + if in_channels != out_channels: + self.conv_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) + else: + self.conv_shortcut = None + + def forward(self, x): + # Apply shortcut connection + residual = x + + # First normalization and activation + x = self.norm1(x) + x = self.nonlinearity(x) + + x = self.conv1(x) + x = self.norm2(x) + x = self.nonlinearity(x) + x = self.conv2(x) + + if self.conv_shortcut is not None: + x = self.conv_shortcut(x) + # Add residual connection + return x + residual + + +class HunyuanImageAttentionBlock(nn.Module): + r""" + Self-attention with a single head. + + Args: + in_channels (int): The number of channels in the input tensor. + """ + + def __init__(self, in_channels: int): + super().__init__() + + # layers + self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + self.to_q = nn.Conv2d(in_channels, in_channels, 1) + self.to_k = nn.Conv2d(in_channels, in_channels, 1) + self.to_v = nn.Conv2d(in_channels, in_channels, 1) + self.proj = nn.Conv2d(in_channels, in_channels, 1) + + def forward(self, x): + identity = x + x = self.norm(x) + + # compute query, key, value + query = self.to_q(x) + key = self.to_k(x) + value = self.to_v(x) + + batch_size, channels, height, width = query.shape + query = query.permute(0, 2, 3, 1).reshape(batch_size, height * width, channels).contiguous() + key = key.permute(0, 2, 3, 1).reshape(batch_size, height * width, channels).contiguous() + value = value.permute(0, 2, 3, 1).reshape(batch_size, height * width, channels).contiguous() + + # apply attention + x = F.scaled_dot_product_attention(query, key, value) + + x = x.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2) + # output projection + x = self.proj(x) + + return x + identity + + +class HunyuanImageDownsample(nn.Module): + """ + Downsampling block for spatial reduction. + + Args: + in_channels (int): Number of input channels. + out_channels (int): Number of output channels. + """ + + def __init__(self, in_channels: int, out_channels: int): + super().__init__() + factor = 4 + if out_channels % factor != 0: + raise ValueError(f"out_channels % factor != 0: {out_channels % factor}") + + self.conv = nn.Conv2d(in_channels, out_channels // factor, kernel_size=3, stride=1, padding=1) + self.group_size = factor * in_channels // out_channels + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h = self.conv(x) + + B, C, H, W = h.shape + h = h.reshape(B, C, H // 2, 2, W // 2, 2) + h = h.permute(0, 3, 5, 1, 2, 4) # b, r1, r2, c, h, w + h = h.reshape(B, 4 * C, H // 2, W // 2) + + B, C, H, W = x.shape + shortcut = x.reshape(B, C, H // 2, 2, W // 2, 2) + shortcut = shortcut.permute(0, 3, 5, 1, 2, 4) # b, r1, r2, c, h, w + shortcut = shortcut.reshape(B, 4 * C, H // 2, W // 2) + + B, C, H, W = shortcut.shape + shortcut = shortcut.view(B, h.shape[1], self.group_size, H, W).mean(dim=2) + return h + shortcut + + +class HunyuanImageUpsample(nn.Module): + """ + Upsampling block for spatial expansion. + + Args: + in_channels (int): Number of input channels. + out_channels (int): Number of output channels. + """ + + def __init__(self, in_channels: int, out_channels: int): + super().__init__() + factor = 4 + self.conv = nn.Conv2d(in_channels, out_channels * factor, kernel_size=3, stride=1, padding=1) + self.repeats = factor * out_channels // in_channels + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h = self.conv(x) + + B, C, H, W = h.shape + h = h.reshape(B, 2, 2, C // 4, H, W) # b, r1, r2, c, h, w + h = h.permute(0, 3, 4, 1, 5, 2) # b, c, h, r1, w, r2 + h = h.reshape(B, C // 4, H * 2, W * 2) + + shortcut = x.repeat_interleave(repeats=self.repeats, dim=1) + + B, C, H, W = shortcut.shape + shortcut = shortcut.reshape(B, 2, 2, C // 4, H, W) # b, r1, r2, c, h, w + shortcut = shortcut.permute(0, 3, 4, 1, 5, 2) # b, c, h, r1, w, r2 + shortcut = shortcut.reshape(B, C // 4, H * 2, W * 2) + return h + shortcut + + +class HunyuanImageMidBlock(nn.Module): + """ + Middle block for HunyuanImageVAE encoder and decoder. + + Args: + in_channels (int): Number of input channels. + num_layers (int): Number of layers. + """ + + def __init__(self, in_channels: int, num_layers: int = 1): + super().__init__() + + resnets = [HunyuanImageResnetBlock(in_channels=in_channels, out_channels=in_channels)] + + attentions = [] + for _ in range(num_layers): + attentions.append(HunyuanImageAttentionBlock(in_channels)) + resnets.append(HunyuanImageResnetBlock(in_channels=in_channels, out_channels=in_channels)) + + self.resnets = nn.ModuleList(resnets) + self.attentions = nn.ModuleList(attentions) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.resnets[0](x) + + for attn, resnet in zip(self.attentions, self.resnets[1:]): + x = attn(x) + x = resnet(x) + + return x + + +class HunyuanImageEncoder2D(nn.Module): + r""" + Encoder network that compresses input to latent representation. + + Args: + in_channels (int): Number of input channels. + z_channels (int): Number of latent channels. + block_out_channels (list of int): Output channels for each block. + num_res_blocks (int): Number of residual blocks per block. + spatial_compression_ratio (int): Spatial downsampling factor. + non_linearity (str): Type of non-linearity to use. Default is "silu". + downsample_match_channel (bool): Whether to match channels during downsampling. + """ + + def __init__( + self, + in_channels: int, + z_channels: int, + block_out_channels: Tuple[int, ...], + num_res_blocks: int, + spatial_compression_ratio: int, + non_linearity: str = "silu", + downsample_match_channel: bool = True, + ): + super().__init__() + if block_out_channels[-1] % (2 * z_channels) != 0: + raise ValueError( + f"block_out_channels[-1 has to be divisible by 2 * out_channels, you have block_out_channels = {block_out_channels[-1]} and out_channels = {z_channels}" + ) + + self.in_channels = in_channels + self.z_channels = z_channels + self.block_out_channels = block_out_channels + self.num_res_blocks = num_res_blocks + self.spatial_compression_ratio = spatial_compression_ratio + + self.group_size = block_out_channels[-1] // (2 * z_channels) + self.nonlinearity = get_activation(non_linearity) + + # init block + self.conv_in = nn.Conv2d(in_channels, block_out_channels[0], kernel_size=3, stride=1, padding=1) + + # downsample blocks + self.down_blocks = nn.ModuleList([]) + + block_in_channel = block_out_channels[0] + for i in range(len(block_out_channels)): + block_out_channel = block_out_channels[i] + # residual blocks + for _ in range(num_res_blocks): + self.down_blocks.append( + HunyuanImageResnetBlock(in_channels=block_in_channel, out_channels=block_out_channel) + ) + block_in_channel = block_out_channel + + # downsample block + if i < np.log2(spatial_compression_ratio) and i != len(block_out_channels) - 1: + if downsample_match_channel: + block_out_channel = block_out_channels[i + 1] + self.down_blocks.append( + HunyuanImageDownsample(in_channels=block_in_channel, out_channels=block_out_channel) + ) + block_in_channel = block_out_channel + + # middle blocks + self.mid_block = HunyuanImageMidBlock(in_channels=block_out_channels[-1], num_layers=1) + + # output blocks + # Output layers + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_out_channels[-1], eps=1e-6, affine=True) + self.conv_out = nn.Conv2d(block_out_channels[-1], 2 * z_channels, kernel_size=3, stride=1, padding=1) + + self.gradient_checkpointing = False + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.conv_in(x) + + ## downsamples + for down_block in self.down_blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + x = self._gradient_checkpointing_func(down_block, x) + else: + x = down_block(x) + + ## middle + if torch.is_grad_enabled() and self.gradient_checkpointing: + x = self._gradient_checkpointing_func(self.mid_block, x) + else: + x = self.mid_block(x) + + ## head + B, C, H, W = x.shape + residual = x.view(B, C // self.group_size, self.group_size, H, W).mean(dim=2) + + x = self.norm_out(x) + x = self.nonlinearity(x) + x = self.conv_out(x) + return x + residual + + +class HunyuanImageDecoder2D(nn.Module): + r""" + Decoder network that reconstructs output from latent representation. + + Args: + z_channels : int + Number of latent channels. + out_channels : int + Number of output channels. + block_out_channels : Tuple[int, ...] + Output channels for each block. + num_res_blocks : int + Number of residual blocks per block. + spatial_compression_ratio : int + Spatial upsampling factor. + upsample_match_channel : bool + Whether to match channels during upsampling. + non_linearity (str): Type of non-linearity to use. Default is "silu". + """ + + def __init__( + self, + z_channels: int, + out_channels: int, + block_out_channels: Tuple[int, ...], + num_res_blocks: int, + spatial_compression_ratio: int, + upsample_match_channel: bool = True, + non_linearity: str = "silu", + ): + super().__init__() + if block_out_channels[0] % z_channels != 0: + raise ValueError( + f"block_out_channels[0] should be divisible by z_channels but has block_out_channels[0] = {block_out_channels[0]} and z_channels = {z_channels}" + ) + + self.z_channels = z_channels + self.block_out_channels = block_out_channels + self.num_res_blocks = num_res_blocks + self.repeat = block_out_channels[0] // z_channels + self.spatial_compression_ratio = spatial_compression_ratio + self.nonlinearity = get_activation(non_linearity) + + self.conv_in = nn.Conv2d(z_channels, block_out_channels[0], kernel_size=3, stride=1, padding=1) + + # Middle blocks with attention + self.mid_block = HunyuanImageMidBlock(in_channels=block_out_channels[0], num_layers=1) + + # Upsampling blocks + block_in_channel = block_out_channels[0] + self.up_blocks = nn.ModuleList() + for i in range(len(block_out_channels)): + block_out_channel = block_out_channels[i] + for _ in range(self.num_res_blocks + 1): + self.up_blocks.append( + HunyuanImageResnetBlock(in_channels=block_in_channel, out_channels=block_out_channel) + ) + block_in_channel = block_out_channel + + if i < np.log2(spatial_compression_ratio) and i != len(block_out_channels) - 1: + if upsample_match_channel: + block_out_channel = block_out_channels[i + 1] + self.up_blocks.append(HunyuanImageUpsample(block_in_channel, block_out_channel)) + block_in_channel = block_out_channel + + # Output layers + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_out_channels[-1], eps=1e-6, affine=True) + self.conv_out = nn.Conv2d(block_out_channels[-1], out_channels, kernel_size=3, stride=1, padding=1) + + self.gradient_checkpointing = False + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h = self.conv_in(x) + x.repeat_interleave(repeats=self.repeat, dim=1) + + if torch.is_grad_enabled() and self.gradient_checkpointing: + h = self._gradient_checkpointing_func(self.mid_block, h) + else: + h = self.mid_block(h) + + for up_block in self.up_blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + h = self._gradient_checkpointing_func(up_block, h) + else: + h = up_block(h) + h = self.norm_out(h) + h = self.nonlinearity(h) + h = self.conv_out(h) + return h + + +class AutoencoderKLHunyuanImage(ModelMixin, ConfigMixin, FromOriginalModelMixin): + r""" + A VAE model for 2D images with spatial tiling support. + + This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented + for all models (such as downloading or saving). + """ + + _supports_gradient_checkpointing = False + + # fmt: off + @register_to_config + def __init__( + self, + in_channels: int, + out_channels: int, + latent_channels: int, + block_out_channels: Tuple[int, ...], + layers_per_block: int, + spatial_compression_ratio: int, + sample_size: int, + scaling_factor: float = None, + downsample_match_channel: bool = True, + upsample_match_channel: bool = True, + ) -> None: + # fmt: on + super().__init__() + + self.encoder = HunyuanImageEncoder2D( + in_channels=in_channels, + z_channels=latent_channels, + block_out_channels=block_out_channels, + num_res_blocks=layers_per_block, + spatial_compression_ratio=spatial_compression_ratio, + downsample_match_channel=downsample_match_channel, + ) + + self.decoder = HunyuanImageDecoder2D( + z_channels=latent_channels, + out_channels=out_channels, + block_out_channels=list(reversed(block_out_channels)), + num_res_blocks=layers_per_block, + spatial_compression_ratio=spatial_compression_ratio, + upsample_match_channel=upsample_match_channel, + ) + + # Tiling and slicing configuration + self.use_slicing = False + self.use_tiling = False + + # Tiling parameters + self.tile_sample_min_size = sample_size + self.tile_latent_min_size = sample_size // spatial_compression_ratio + self.tile_overlap_factor = 0.25 + + def enable_tiling( + self, + tile_sample_min_size: Optional[int] = None, + tile_overlap_factor: Optional[float] = None, + ) -> None: + r""" + Enable spatial tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles + to compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to + allow processing larger images. + + Args: + tile_sample_min_size (`int`, *optional*): + The minimum size required for a sample to be separated into tiles across the spatial dimension. + tile_overlap_factor (`float`, *optional*): + The overlap factor required for a latent to be separated into tiles across the spatial dimension. + """ + self.use_tiling = True + self.tile_sample_min_size = tile_sample_min_size or self.tile_sample_min_size + self.tile_overlap_factor = tile_overlap_factor or self.tile_overlap_factor + self.tile_latent_min_size = self.tile_sample_min_size // self.config.spatial_compression_ratio + + def disable_tiling(self) -> None: + r""" + Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing + decoding in one step. + """ + self.use_tiling = False + + def enable_slicing(self) -> None: + r""" + Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to + compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. + """ + self.use_slicing = True + + def disable_slicing(self) -> None: + r""" + Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing + decoding in one step. + """ + self.use_slicing = False + + def _encode(self, x: torch.Tensor): + + batch_size, num_channels, height, width = x.shape + + if self.use_tiling and (width > self.tile_sample_min_size or height > self.tile_sample_min_size): + return self.tiled_encode(x) + + enc = self.encoder(x) + + return enc + + @apply_forward_hook + def encode( + self, x: torch.Tensor, return_dict: bool = True + ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: + r""" + Encode a batch of images into latents. + + Args: + x (`torch.Tensor`): Input batch of images. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. + + Returns: + The latent representations of the encoded videos. If `return_dict` is True, a + [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned. + """ + if self.use_slicing and x.shape[0] > 1: + encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)] + h = torch.cat(encoded_slices) + else: + h = self._encode(x) + posterior = DiagonalGaussianDistribution(h) + + if not return_dict: + return (posterior,) + return AutoencoderKLOutput(latent_dist=posterior) + + def _decode(self, z: torch.Tensor, return_dict: bool = True): + + batch_size, num_channels, height, width = z.shape + + if self.use_tiling and (width > self.tile_latent_min_size or height > self.tile_latent_min_size): + return self.tiled_decode(z, return_dict=return_dict) + + dec = self.decoder(z) + + if not return_dict: + return (dec,) + + return DecoderOutput(sample=dec) + + @apply_forward_hook + def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + r""" + Decode a batch of images. + + Args: + z (`torch.Tensor`): Input batch of latent vectors. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is + returned. + """ + if self.use_slicing and z.shape[0] > 1: + decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)] + decoded = torch.cat(decoded_slices) + else: + decoded = self._decode(z).sample + + if not return_dict: + return (decoded,) + return DecoderOutput(sample=decoded) + + + def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-2], b.shape[-2], blend_extent) + for y in range(blend_extent): + b[:, :, y, :] = a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * ( + y / blend_extent + ) + return b + + def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-1], b.shape[-1], blend_extent) + for x in range(blend_extent): + b[:, :, :, x] = a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * ( + x / blend_extent + ) + return b + + def tiled_encode(self, x: torch.Tensor) -> torch.Tensor: + """ + Encode input using spatial tiling strategy. + + Args: + x (`torch.Tensor`): Input tensor of shape (B, C, T, H, W). + + Returns: + `torch.Tensor`: + The latent representation of the encoded images. + """ + _, _, _, height, width = x.shape + overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor) + row_limit = self.tile_latent_min_size - blend_extent + + rows = [] + for i in range(0, height, overlap_size): + row = [] + for j in range(0, width, overlap_size): + tile = x[:, :, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size] + tile = self.encoder(tile) + row.append(tile) + rows.append(row) + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + + moments = torch.cat(result_rows, dim=-2) + + return moments + + def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + """ + Decode latent using spatial tiling strategy. + + Args: + z (`torch.Tensor`): Latent tensor of shape (B, C, H, W). + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is + returned. + """ + _, _, height, width = z.shape + overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor) + row_limit = self.tile_sample_min_size - blend_extent + + rows = [] + for i in range(0, height, overlap_size): + row = [] + for j in range(0, width, overlap_size): + tile = z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size] + decoded = self.decoder(tile) + row.append(decoded) + rows.append(row) + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + + dec = torch.cat(result_rows, dim=-2) + if not return_dict: + return (dec,) + return DecoderOutput(sample=dec) + + + def forward( + self, + sample: torch.Tensor, + sample_posterior: bool = False, + return_dict: bool = True, + generator: Optional[torch.Generator] = None, + ) -> Union[DecoderOutput, torch.Tensor]: + """ + Args: + sample (`torch.Tensor`): Input sample. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`DecoderOutput`] instead of a plain tuple. + """ + posterior = self.encode(sample).latent_dist + if sample_posterior: + z = posterior.sample(generator=generator) + else: + z = posterior.mode() + dec = self.decode(z, return_dict=return_dict) + + return dec diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage_refiner.py b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage_refiner.py new file mode 100644 index 0000000000..af40c7a6cb --- /dev/null +++ b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage_refiner.py @@ -0,0 +1,934 @@ +# Copyright 2025 The Hunyuan Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Optional, Tuple, Union + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint + +from ...configuration_utils import ConfigMixin, register_to_config +from ...utils import logging +from ...utils.accelerate_utils import apply_forward_hook +from ..activations import get_activation +from ..modeling_outputs import AutoencoderKLOutput +from ..modeling_utils import ModelMixin +from .vae import DecoderOutput, DiagonalGaussianDistribution + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class HunyuanImageRefinerCausalConv3d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: Union[int, Tuple[int, int, int]] = 3, + stride: Union[int, Tuple[int, int, int]] = 1, + padding: Union[int, Tuple[int, int, int]] = 0, + dilation: Union[int, Tuple[int, int, int]] = 1, + bias: bool = True, + pad_mode: str = "replicate", + ) -> None: + super().__init__() + + kernel_size = (kernel_size, kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size + + self.pad_mode = pad_mode + self.time_causal_padding = ( + kernel_size[0] // 2, + kernel_size[0] // 2, + kernel_size[1] // 2, + kernel_size[1] // 2, + kernel_size[2] - 1, + 0, + ) + + self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = F.pad(hidden_states, self.time_causal_padding, mode=self.pad_mode) + return self.conv(hidden_states) + + +class HunyuanImageRefinerRMS_norm(nn.Module): + r""" + A custom RMS normalization layer. + + Args: + dim (int): The number of dimensions to normalize over. + channel_first (bool, optional): Whether the input tensor has channels as the first dimension. + Default is True. + images (bool, optional): Whether the input represents image data. Default is True. + bias (bool, optional): Whether to include a learnable bias term. Default is False. + """ + + def __init__(self, dim: int, channel_first: bool = True, images: bool = True, bias: bool = False) -> None: + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 + + def forward(self, x): + return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias + + +class HunyuanImageRefinerAttnBlock(nn.Module): + def __init__(self, in_channels: int): + super().__init__() + self.in_channels = in_channels + + self.norm = HunyuanImageRefinerRMS_norm(in_channels, images=False) + + self.to_q = nn.Conv3d(in_channels, in_channels, kernel_size=1) + self.to_k = nn.Conv3d(in_channels, in_channels, kernel_size=1) + self.to_v = nn.Conv3d(in_channels, in_channels, kernel_size=1) + self.proj_out = nn.Conv3d(in_channels, in_channels, kernel_size=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + identity = x + + x = self.norm(x) + + query = self.to_q(x) + key = self.to_k(x) + value = self.to_v(x) + + batch_size, channels, frames, height, width = query.shape + + query = query.reshape(batch_size, channels, frames * height * width).permute(0, 2, 1).unsqueeze(1).contiguous() + key = key.reshape(batch_size, channels, frames * height * width).permute(0, 2, 1).unsqueeze(1).contiguous() + value = value.reshape(batch_size, channels, frames * height * width).permute(0, 2, 1).unsqueeze(1).contiguous() + + x = nn.functional.scaled_dot_product_attention(query, key, value, attn_mask=None) + + # batch_size, 1, frames * height * width, channels + + x = x.squeeze(1).reshape(batch_size, frames, height, width, channels).permute(0, 4, 1, 2, 3) + x = self.proj_out(x) + + return x + identity + + +class HunyuanImageRefinerUpsampleDCAE(nn.Module): + def __init__(self, in_channels: int, out_channels: int, add_temporal_upsample: bool = True): + super().__init__() + factor = 2 * 2 * 2 if add_temporal_upsample else 1 * 2 * 2 + self.conv = HunyuanImageRefinerCausalConv3d(in_channels, out_channels * factor, kernel_size=3) + + self.add_temporal_upsample = add_temporal_upsample + self.repeats = factor * out_channels // in_channels + + @staticmethod + def _dcae_upsample_rearrange(tensor, r1=1, r2=2, r3=2): + """ + Convert (b, r1*r2*r3*c, f, h, w) -> (b, c, r1*f, r2*h, r3*w) + + Args: + tensor: Input tensor of shape (b, r1*r2*r3*c, f, h, w) + r1: temporal upsampling factor + r2: height upsampling factor + r3: width upsampling factor + """ + b, packed_c, f, h, w = tensor.shape + factor = r1 * r2 * r3 + c = packed_c // factor + + tensor = tensor.view(b, r1, r2, r3, c, f, h, w) + tensor = tensor.permute(0, 4, 5, 1, 6, 2, 7, 3) + return tensor.reshape(b, c, f * r1, h * r2, w * r3) + + def forward(self, x: torch.Tensor): + r1 = 2 if self.add_temporal_upsample else 1 + h = self.conv(x) + if self.add_temporal_upsample: + h = self._dcae_upsample_rearrange(h, r1=1, r2=2, r3=2) + h = h[:, : h.shape[1] // 2] + + # shortcut computation + shortcut = self._dcae_upsample_rearrange(x, r1=1, r2=2, r3=2) + shortcut = shortcut.repeat_interleave(repeats=self.repeats // 2, dim=1) + + else: + h = self._dcae_upsample_rearrange(h, r1=r1, r2=2, r3=2) + shortcut = x.repeat_interleave(repeats=self.repeats, dim=1) + shortcut = self._dcae_upsample_rearrange(shortcut, r1=r1, r2=2, r3=2) + return h + shortcut + + +class HunyuanImageRefinerDownsampleDCAE(nn.Module): + def __init__(self, in_channels: int, out_channels: int, add_temporal_downsample: bool = True): + super().__init__() + factor = 2 * 2 * 2 if add_temporal_downsample else 1 * 2 * 2 + assert out_channels % factor == 0 + # self.conv = Conv3d(in_channels, out_channels // factor, kernel_size=3, stride=1, padding=1) + self.conv = HunyuanImageRefinerCausalConv3d(in_channels, out_channels // factor, kernel_size=3) + + self.add_temporal_downsample = add_temporal_downsample + self.group_size = factor * in_channels // out_channels + + @staticmethod + def _dcae_downsample_rearrange(tensor, r1=1, r2=2, r3=2): + """ + Convert (b, c, r1*f, r2*h, r3*w) -> (b, r1*r2*r3*c, f, h, w) + + This packs spatial/temporal dimensions into channels (opposite of upsample) + """ + b, c, packed_f, packed_h, packed_w = tensor.shape + f, h, w = packed_f // r1, packed_h // r2, packed_w // r3 + + tensor = tensor.view(b, c, f, r1, h, r2, w, r3) + tensor = tensor.permute(0, 3, 5, 7, 1, 2, 4, 6) + return tensor.reshape(b, r1 * r2 * r3 * c, f, h, w) + + def forward(self, x: torch.Tensor): + r1 = 2 if self.add_temporal_downsample else 1 + h = self.conv(x) + if self.add_temporal_downsample: + # h = rearrange(h, "b c f (h r2) (w r3) -> b (r2 r3 c) f h w", r2=2, r3=2) + h = self._dcae_downsample_rearrange(h, r1=1, r2=2, r3=2) + h = torch.cat([h, h], dim=1) + # shortcut computation + # shortcut = rearrange(x, "b c f (h r2) (w r3) -> b (r2 r3 c) f h w", r2=2, r3=2) + shortcut = self._dcae_downsample_rearrange(x, r1=1, r2=2, r3=2) + B, C, T, H, W = shortcut.shape + shortcut = shortcut.view(B, h.shape[1], self.group_size // 2, T, H, W).mean(dim=2) + else: + # h = rearrange(h, "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=r1, r2=2, r3=2) + h = self._dcae_downsample_rearrange(h, r1=r1, r2=2, r3=2) + # shortcut = rearrange(x, "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=r1, r2=2, r3=2) + shortcut = self._dcae_downsample_rearrange(x, r1=r1, r2=2, r3=2) + B, C, T, H, W = shortcut.shape + shortcut = shortcut.view(B, h.shape[1], self.group_size, T, H, W).mean(dim=2) + + return h + shortcut + + +class HunyuanImageRefinerResnetBlock(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: Optional[int] = None, + non_linearity: str = "swish", + ) -> None: + super().__init__() + out_channels = out_channels or in_channels + + self.nonlinearity = get_activation(non_linearity) + + self.norm1 = HunyuanImageRefinerRMS_norm(in_channels, images=False) + self.conv1 = HunyuanImageRefinerCausalConv3d(in_channels, out_channels, kernel_size=3) + + self.norm2 = HunyuanImageRefinerRMS_norm(out_channels, images=False) + self.conv2 = HunyuanImageRefinerCausalConv3d(out_channels, out_channels, kernel_size=3) + + self.conv_shortcut = None + if in_channels != out_channels: + self.conv_shortcut = nn.Conv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + residual = hidden_states + + hidden_states = self.norm1(hidden_states) + hidden_states = self.nonlinearity(hidden_states) + hidden_states = self.conv1(hidden_states) + + hidden_states = self.norm2(hidden_states) + hidden_states = self.nonlinearity(hidden_states) + hidden_states = self.conv2(hidden_states) + + if self.conv_shortcut is not None: + residual = self.conv_shortcut(residual) + + return hidden_states + residual + + +class HunyuanImageRefinerMidBlock(nn.Module): + def __init__( + self, + in_channels: int, + num_layers: int = 1, + add_attention: bool = True, + ) -> None: + super().__init__() + self.add_attention = add_attention + + # There is always at least one resnet + resnets = [ + HunyuanImageRefinerResnetBlock( + in_channels=in_channels, + out_channels=in_channels, + ) + ] + attentions = [] + + for _ in range(num_layers): + if self.add_attention: + attentions.append(HunyuanImageRefinerAttnBlock(in_channels)) + else: + attentions.append(None) + + resnets.append( + HunyuanImageRefinerResnetBlock( + in_channels=in_channels, + out_channels=in_channels, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + self.gradient_checkpointing = False + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.resnets[0](hidden_states) + + for attn, resnet in zip(self.attentions, self.resnets[1:]): + if attn is not None: + hidden_states = attn(hidden_states) + hidden_states = resnet(hidden_states) + + return hidden_states + + +class HunyuanImageRefinerDownBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + num_layers: int = 1, + downsample_out_channels: Optional[int] = None, + add_temporal_downsample: int = True, + ) -> None: + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + HunyuanImageRefinerResnetBlock( + in_channels=in_channels, + out_channels=out_channels, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if downsample_out_channels is not None: + self.downsamplers = nn.ModuleList( + [ + HunyuanImageRefinerDownsampleDCAE( + out_channels, + out_channels=downsample_out_channels, + add_temporal_downsample=add_temporal_downsample, + ) + ] + ) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + for resnet in self.resnets: + hidden_states = resnet(hidden_states) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + + return hidden_states + + +class HunyuanImageRefinerUpBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + num_layers: int = 1, + upsample_out_channels: Optional[int] = None, + add_temporal_upsample: bool = True, + ) -> None: + super().__init__() + resnets = [] + + for i in range(num_layers): + input_channels = in_channels if i == 0 else out_channels + + resnets.append( + HunyuanImageRefinerResnetBlock( + in_channels=input_channels, + out_channels=out_channels, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if upsample_out_channels is not None: + self.upsamplers = nn.ModuleList( + [ + HunyuanImageRefinerUpsampleDCAE( + out_channels, + out_channels=upsample_out_channels, + add_temporal_upsample=add_temporal_upsample, + ) + ] + ) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if torch.is_grad_enabled() and self.gradient_checkpointing: + for resnet in self.resnets: + hidden_states = self._gradient_checkpointing_func(resnet, hidden_states) + + else: + for resnet in self.resnets: + hidden_states = resnet(hidden_states) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states) + + return hidden_states + + +class HunyuanImageRefinerEncoder3D(nn.Module): + r""" + 3D vae encoder for HunyuanImageRefiner. + """ + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 64, + block_out_channels: Tuple[int, ...] = (128, 256, 512, 1024, 1024), + layers_per_block: int = 2, + temporal_compression_ratio: int = 4, + spatial_compression_ratio: int = 16, + downsample_match_channel: bool = True, + ) -> None: + super().__init__() + + self.in_channels = in_channels + self.out_channels = out_channels + self.group_size = block_out_channels[-1] // self.out_channels + + self.conv_in = HunyuanImageRefinerCausalConv3d(in_channels, block_out_channels[0], kernel_size=3) + self.mid_block = None + self.down_blocks = nn.ModuleList([]) + + input_channel = block_out_channels[0] + for i in range(len(block_out_channels)): + add_spatial_downsample = i < np.log2(spatial_compression_ratio) + output_channel = block_out_channels[i] + if not add_spatial_downsample: + down_block = HunyuanImageRefinerDownBlock3D( + num_layers=layers_per_block, + in_channels=input_channel, + out_channels=output_channel, + downsample_out_channels=None, + add_temporal_downsample=False, + ) + input_channel = output_channel + else: + add_temporal_downsample = i >= np.log2(spatial_compression_ratio // temporal_compression_ratio) + downsample_out_channels = block_out_channels[i + 1] if downsample_match_channel else output_channel + down_block = HunyuanImageRefinerDownBlock3D( + num_layers=layers_per_block, + in_channels=input_channel, + out_channels=output_channel, + downsample_out_channels=downsample_out_channels, + add_temporal_downsample=add_temporal_downsample, + ) + input_channel = downsample_out_channels + + self.down_blocks.append(down_block) + + self.mid_block = HunyuanImageRefinerMidBlock(in_channels=block_out_channels[-1]) + + self.norm_out = HunyuanImageRefinerRMS_norm(block_out_channels[-1], images=False) + self.conv_act = nn.SiLU() + self.conv_out = HunyuanImageRefinerCausalConv3d(block_out_channels[-1], out_channels, kernel_size=3) + + self.gradient_checkpointing = False + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.conv_in(hidden_states) + + if torch.is_grad_enabled() and self.gradient_checkpointing: + for down_block in self.down_blocks: + hidden_states = self._gradient_checkpointing_func(down_block, hidden_states) + + hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states) + else: + for down_block in self.down_blocks: + hidden_states = down_block(hidden_states) + + hidden_states = self.mid_block(hidden_states) + + # short_cut = rearrange(hidden_states, "b (c r) f h w -> b c r f h w", r=self.group_size).mean(dim=2) + batch_size, _, frame, height, width = hidden_states.shape + short_cut = hidden_states.view(batch_size, -1, self.group_size, frame, height, width).mean(dim=2) + + hidden_states = self.norm_out(hidden_states) + hidden_states = self.conv_act(hidden_states) + hidden_states = self.conv_out(hidden_states) + + hidden_states += short_cut + + return hidden_states + + +class HunyuanImageRefinerDecoder3D(nn.Module): + r""" + Causal decoder for 3D video-like data used for HunyuanImage-2.1 Refiner. + """ + + def __init__( + self, + in_channels: int = 32, + out_channels: int = 3, + block_out_channels: Tuple[int, ...] = (1024, 1024, 512, 256, 128), + layers_per_block: int = 2, + spatial_compression_ratio: int = 16, + temporal_compression_ratio: int = 4, + upsample_match_channel: bool = True, + ): + super().__init__() + self.layers_per_block = layers_per_block + self.in_channels = in_channels + self.out_channels = out_channels + self.repeat = block_out_channels[0] // self.in_channels + + self.conv_in = HunyuanImageRefinerCausalConv3d(self.in_channels, block_out_channels[0], kernel_size=3) + self.up_blocks = nn.ModuleList([]) + + # mid + self.mid_block = HunyuanImageRefinerMidBlock(in_channels=block_out_channels[0]) + + # up + input_channel = block_out_channels[0] + for i in range(len(block_out_channels)): + output_channel = block_out_channels[i] + + add_spatial_upsample = i < np.log2(spatial_compression_ratio) + add_temporal_upsample = i < np.log2(temporal_compression_ratio) + if add_spatial_upsample or add_temporal_upsample: + upsample_out_channels = block_out_channels[i + 1] if upsample_match_channel else output_channel + up_block = HunyuanImageRefinerUpBlock3D( + num_layers=self.layers_per_block + 1, + in_channels=input_channel, + out_channels=output_channel, + upsample_out_channels=upsample_out_channels, + add_temporal_upsample=add_temporal_upsample, + ) + input_channel = upsample_out_channels + else: + up_block = HunyuanImageRefinerUpBlock3D( + num_layers=self.layers_per_block + 1, + in_channels=input_channel, + out_channels=output_channel, + upsample_out_channels=None, + add_temporal_upsample=False, + ) + input_channel = output_channel + + self.up_blocks.append(up_block) + + # out + self.norm_out = HunyuanImageRefinerRMS_norm(block_out_channels[-1], images=False) + self.conv_act = nn.SiLU() + self.conv_out = HunyuanImageRefinerCausalConv3d(block_out_channels[-1], out_channels, kernel_size=3) + + self.gradient_checkpointing = False + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.conv_in(hidden_states) + hidden_states.repeat_interleave(repeats=self.repeat, dim=1) + + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states) + + for up_block in self.up_blocks: + hidden_states = self._gradient_checkpointing_func(up_block, hidden_states) + else: + hidden_states = self.mid_block(hidden_states) + + for up_block in self.up_blocks: + hidden_states = up_block(hidden_states) + + # post-process + hidden_states = self.norm_out(hidden_states) + hidden_states = self.conv_act(hidden_states) + hidden_states = self.conv_out(hidden_states) + return hidden_states + + +class AutoencoderKLHunyuanImageRefiner(ModelMixin, ConfigMixin): + r""" + A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. Used for + HunyuanImage-2.1 Refiner. + + This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented + for all models (such as downloading or saving). + """ + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + latent_channels: int = 32, + block_out_channels: Tuple[int] = (128, 256, 512, 1024, 1024), + layers_per_block: int = 2, + spatial_compression_ratio: int = 16, + temporal_compression_ratio: int = 4, + downsample_match_channel: bool = True, + upsample_match_channel: bool = True, + scaling_factor: float = 1.03682, + ) -> None: + super().__init__() + + self.encoder = HunyuanImageRefinerEncoder3D( + in_channels=in_channels, + out_channels=latent_channels * 2, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + temporal_compression_ratio=temporal_compression_ratio, + spatial_compression_ratio=spatial_compression_ratio, + downsample_match_channel=downsample_match_channel, + ) + + self.decoder = HunyuanImageRefinerDecoder3D( + in_channels=latent_channels, + out_channels=out_channels, + block_out_channels=list(reversed(block_out_channels)), + layers_per_block=layers_per_block, + temporal_compression_ratio=temporal_compression_ratio, + spatial_compression_ratio=spatial_compression_ratio, + upsample_match_channel=upsample_match_channel, + ) + + self.spatial_compression_ratio = spatial_compression_ratio + self.temporal_compression_ratio = temporal_compression_ratio + + # When decoding a batch of video latents at a time, one can save memory by slicing across the batch dimension + # to perform decoding of a single video latent at a time. + self.use_slicing = False + + # When decoding spatially large video latents, the memory requirement is very high. By breaking the video latent + # frames spatially into smaller tiles and performing multiple forward passes for decoding, and then blending the + # intermediate tiles together, the memory requirement can be lowered. + self.use_tiling = False + + # The minimal tile height and width for spatial tiling to be used + self.tile_sample_min_height = 256 + self.tile_sample_min_width = 256 + + # The minimal distance between two spatial tiles + self.tile_sample_stride_height = 192 + self.tile_sample_stride_width = 192 + + self.tile_overlap_factor = 0.25 + + def enable_tiling( + self, + tile_sample_min_height: Optional[int] = None, + tile_sample_min_width: Optional[int] = None, + tile_sample_stride_height: Optional[float] = None, + tile_sample_stride_width: Optional[float] = None, + tile_overlap_factor: Optional[float] = None, + ) -> None: + r""" + Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to + compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow + processing larger images. + + Args: + tile_sample_min_height (`int`, *optional*): + The minimum height required for a sample to be separated into tiles across the height dimension. + tile_sample_min_width (`int`, *optional*): + The minimum width required for a sample to be separated into tiles across the width dimension. + tile_sample_stride_height (`int`, *optional*): + The minimum amount of overlap between two consecutive vertical tiles. This is to ensure that there are + no tiling artifacts produced across the height dimension. + tile_sample_stride_width (`int`, *optional*): + The stride between two consecutive horizontal tiles. This is to ensure that there are no tiling + artifacts produced across the width dimension. + """ + self.use_tiling = True + self.tile_sample_min_height = tile_sample_min_height or self.tile_sample_min_height + self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width + self.tile_sample_stride_height = tile_sample_stride_height or self.tile_sample_stride_height + self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width + self.tile_overlap_factor = tile_overlap_factor or self.tile_overlap_factor + + def disable_tiling(self) -> None: + r""" + Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing + decoding in one step. + """ + self.use_tiling = False + + def enable_slicing(self) -> None: + r""" + Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to + compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. + """ + self.use_slicing = True + + def disable_slicing(self) -> None: + r""" + Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing + decoding in one step. + """ + self.use_slicing = False + + def _encode(self, x: torch.Tensor) -> torch.Tensor: + _, _, _, height, width = x.shape + + if self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height): + return self.tiled_encode(x) + + x = self.encoder(x) + return x + + @apply_forward_hook + def encode( + self, x: torch.Tensor, return_dict: bool = True + ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: + r""" + Encode a batch of images into latents. + + Args: + x (`torch.Tensor`): Input batch of images. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. + + Returns: + The latent representations of the encoded videos. If `return_dict` is True, a + [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned. + """ + if self.use_slicing and x.shape[0] > 1: + encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)] + h = torch.cat(encoded_slices) + else: + h = self._encode(x) + + posterior = DiagonalGaussianDistribution(h) + + if not return_dict: + return (posterior,) + return AutoencoderKLOutput(latent_dist=posterior) + + def _decode(self, z: torch.Tensor) -> torch.Tensor: + _, _, _, height, width = z.shape + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + + if self.use_tiling and (width > tile_latent_min_width or height > tile_latent_min_height): + return self.tiled_decode(z) + + dec = self.decoder(z) + + return dec + + @apply_forward_hook + def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + r""" + Decode a batch of images. + + Args: + z (`torch.Tensor`): Input batch of latent vectors. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is + returned. + """ + if self.use_slicing and z.shape[0] > 1: + decoded_slices = [self._decode(z_slice) for z_slice in z.split(1)] + decoded = torch.cat(decoded_slices) + else: + decoded = self._decode(z) + + if not return_dict: + return (decoded,) + + return DecoderOutput(sample=decoded) + + def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-2], b.shape[-2], blend_extent) + for y in range(blend_extent): + b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * ( + y / blend_extent + ) + return b + + def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-1], b.shape[-1], blend_extent) + for x in range(blend_extent): + b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * ( + x / blend_extent + ) + return b + + def blend_t(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-3], b.shape[-3], blend_extent) + for x in range(blend_extent): + b[:, :, x, :, :] = a[:, :, -blend_extent + x, :, :] * (1 - x / blend_extent) + b[:, :, x, :, :] * ( + x / blend_extent + ) + return b + + def tiled_encode(self, x: torch.Tensor) -> torch.Tensor: + r"""Encode a batch of images using a tiled encoder. + + Args: + x (`torch.Tensor`): Input batch of videos. + + Returns: + `torch.Tensor`: + The latent representation of the encoded videos. + """ + _, _, _, height, width = x.shape + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + overlap_height = int(tile_latent_min_height * (1 - self.tile_overlap_factor)) # 256 * (1 - 0.25) = 192 + overlap_width = int(tile_latent_min_width * (1 - self.tile_overlap_factor)) # 256 * (1 - 0.25) = 192 + blend_height = int(tile_latent_min_height * self.tile_overlap_factor) # 8 * 0.25 = 2 + blend_width = int(tile_latent_min_width * self.tile_overlap_factor) # 8 * 0.25 = 2 + row_limit_height = tile_latent_min_height - blend_height # 8 - 2 = 6 + row_limit_width = tile_latent_min_width - blend_width # 8 - 2 = 6 + + rows = [] + for i in range(0, height, overlap_height): + row = [] + for j in range(0, width, overlap_width): + tile = x[ + :, + :, + :, + i : i + self.tile_sample_min_height, + j : j + self.tile_sample_min_width, + ] + tile = self.encoder(tile) + row.append(tile) + rows.append(row) + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_height) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_width) + result_row.append(tile[:, :, :, :row_limit_height, :row_limit_width]) + result_rows.append(torch.cat(result_row, dim=-1)) + moments = torch.cat(result_rows, dim=-2) + + return moments + + def tiled_decode(self, z: torch.Tensor) -> torch.Tensor: + r""" + Decode a batch of images using a tiled decoder. + + Args: + z (`torch.Tensor`): Input batch of latent vectors. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is + returned. + """ + + _, _, _, height, width = z.shape + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + overlap_height = int(tile_latent_min_height * (1 - self.tile_overlap_factor)) # 8 * (1 - 0.25) = 6 + overlap_width = int(tile_latent_min_width * (1 - self.tile_overlap_factor)) # 8 * (1 - 0.25) = 6 + blend_height = int(tile_latent_min_height * self.tile_overlap_factor) # 256 * 0.25 = 64 + blend_width = int(tile_latent_min_width * self.tile_overlap_factor) # 256 * 0.25 = 64 + row_limit_height = tile_latent_min_height - blend_height # 256 - 64 = 192 + row_limit_width = tile_latent_min_width - blend_width # 256 - 64 = 192 + + rows = [] + for i in range(0, height, overlap_height): + row = [] + for j in range(0, width, overlap_width): + tile = z[ + :, + :, + :, + i : i + tile_latent_min_height, + j : j + tile_latent_min_width, + ] + decoded = self.decoder(tile) + row.append(decoded) + rows.append(row) + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_height) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_width) + result_row.append(tile[:, :, :, :row_limit_height, :row_limit_width]) + result_rows.append(torch.cat(result_row, dim=-1)) + dec = torch.cat(result_rows, dim=-2) + + return dec + + def forward( + self, + sample: torch.Tensor, + sample_posterior: bool = False, + return_dict: bool = True, + generator: Optional[torch.Generator] = None, + ) -> Union[DecoderOutput, torch.Tensor]: + r""" + Args: + sample (`torch.Tensor`): Input sample. + sample_posterior (`bool`, *optional*, defaults to `False`): + Whether to sample from the posterior. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`DecoderOutput`] instead of a plain tuple. + """ + x = sample + posterior = self.encode(x).latent_dist + if sample_posterior: + z = posterior.sample(generator=generator) + else: + z = posterior.mode() + dec = self.decode(z, return_dict=return_dict) + return dec diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index d8c3d9b572..6c66131dea 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -27,6 +27,7 @@ if is_torch_available(): from .transformer_hidream_image import HiDreamImageTransformer2DModel from .transformer_hunyuan_video import HunyuanVideoTransformer3DModel from .transformer_hunyuan_video_framepack import HunyuanVideoFramepackTransformer3DModel + from .transformer_hunyuanimage import HunyuanImageTransformer2DModel from .transformer_kandinsky import Kandinsky5Transformer3DModel from .transformer_ltx import LTXVideoTransformer3DModel from .transformer_lumina2 import Lumina2Transformer2DModel diff --git a/src/diffusers/models/transformers/transformer_hunyuanimage.py b/src/diffusers/models/transformers/transformer_hunyuanimage.py new file mode 100644 index 0000000000..7f37bf815b --- /dev/null +++ b/src/diffusers/models/transformers/transformer_hunyuanimage.py @@ -0,0 +1,971 @@ +# Copyright 2025 The Hunyuan Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from diffusers.loaders import FromOriginalModelMixin + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import PeftAdapterMixin +from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers +from ...utils.torch_utils import maybe_allow_in_graph +from ..attention import FeedForward +from ..attention_dispatch import dispatch_attention_fn +from ..attention_processor import Attention, AttentionProcessor +from ..cache_utils import CacheMixin +from ..embeddings import ( + CombinedTimestepTextProjEmbeddings, + TimestepEmbedding, + Timesteps, + get_1d_rotary_pos_embed, +) +from ..modeling_outputs import Transformer2DModelOutput +from ..modeling_utils import ModelMixin +from ..normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class HunyuanImageAttnProcessor: + _attention_backend = None + _parallel_config = None + + def __init__(self): + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError( + "HunyuanImageAttnProcessor requires PyTorch 2.0. To use it, please upgrade PyTorch to 2.0." + ) + + def __call__( + self, + attn: Attention, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + image_rotary_emb: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if attn.add_q_proj is None and encoder_hidden_states is not None: + hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1) + + # 1. QKV projections + query = attn.to_q(hidden_states) + key = attn.to_k(hidden_states) + value = attn.to_v(hidden_states) + + query = query.unflatten(2, (attn.heads, -1)) # batch_size, seq_len, heads, head_dim + key = key.unflatten(2, (attn.heads, -1)) + value = value.unflatten(2, (attn.heads, -1)) + + # 2. QK normalization + if attn.norm_q is not None: + query = attn.norm_q(query) + if attn.norm_k is not None: + key = attn.norm_k(key) + + # 3. Rotational positional embeddings applied to latent stream + if image_rotary_emb is not None: + from ..embeddings import apply_rotary_emb + + if attn.add_q_proj is None and encoder_hidden_states is not None: + query = torch.cat( + [ + apply_rotary_emb( + query[:, : -encoder_hidden_states.shape[1]], image_rotary_emb, sequence_dim=1 + ), + query[:, -encoder_hidden_states.shape[1] :], + ], + dim=1, + ) + key = torch.cat( + [ + apply_rotary_emb(key[:, : -encoder_hidden_states.shape[1]], image_rotary_emb, sequence_dim=1), + key[:, -encoder_hidden_states.shape[1] :], + ], + dim=1, + ) + else: + query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1) + key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1) + + # 4. Encoder condition QKV projection and normalization + if attn.add_q_proj is not None and encoder_hidden_states is not None: + encoder_query = attn.add_q_proj(encoder_hidden_states) + encoder_key = attn.add_k_proj(encoder_hidden_states) + encoder_value = attn.add_v_proj(encoder_hidden_states) + + encoder_query = encoder_query.unflatten(2, (attn.heads, -1)) + encoder_key = encoder_key.unflatten(2, (attn.heads, -1)) + encoder_value = encoder_value.unflatten(2, (attn.heads, -1)) + + if attn.norm_added_q is not None: + encoder_query = attn.norm_added_q(encoder_query) + if attn.norm_added_k is not None: + encoder_key = attn.norm_added_k(encoder_key) + + query = torch.cat([query, encoder_query], dim=1) + key = torch.cat([key, encoder_key], dim=1) + value = torch.cat([value, encoder_value], dim=1) + + # 5. Attention + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.to(query.dtype) + + # 6. Output projection + if encoder_hidden_states is not None: + hidden_states, encoder_hidden_states = ( + hidden_states[:, : -encoder_hidden_states.shape[1]], + hidden_states[:, -encoder_hidden_states.shape[1] :], + ) + + if getattr(attn, "to_out", None) is not None: + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + + if getattr(attn, "to_add_out", None) is not None: + encoder_hidden_states = attn.to_add_out(encoder_hidden_states) + + return hidden_states, encoder_hidden_states + + +class HunyuanImagePatchEmbed(nn.Module): + def __init__( + self, + patch_size: Union[Tuple[int, int], Tuple[int, int, int]] = (16, 16), + in_chans: int = 3, + embed_dim: int = 768, + ) -> None: + super().__init__() + + self.patch_size = patch_size + + if len(patch_size) == 2: + self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) + elif len(patch_size) == 3: + self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) + else: + raise ValueError(f"patch_size must be a tuple of length 2 or 3, got {len(patch_size)}") + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.proj(hidden_states) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + return hidden_states + + +class HunyuanImageByT5TextProjection(nn.Module): + def __init__(self, in_features: int, hidden_size: int, out_features: int): + super().__init__() + self.norm = nn.LayerNorm(in_features) + self.linear_1 = nn.Linear(in_features, hidden_size) + self.linear_2 = nn.Linear(hidden_size, hidden_size) + self.linear_3 = nn.Linear(hidden_size, out_features) + self.act_fn = nn.GELU() + + def forward(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.norm(encoder_hidden_states) + hidden_states = self.linear_1(hidden_states) + hidden_states = self.act_fn(hidden_states) + hidden_states = self.linear_2(hidden_states) + hidden_states = self.act_fn(hidden_states) + hidden_states = self.linear_3(hidden_states) + return hidden_states + + +class HunyuanImageAdaNorm(nn.Module): + def __init__(self, in_features: int, out_features: Optional[int] = None) -> None: + super().__init__() + + out_features = out_features or 2 * in_features + self.linear = nn.Linear(in_features, out_features) + self.nonlinearity = nn.SiLU() + + def forward( + self, temb: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + temb = self.linear(self.nonlinearity(temb)) + gate_msa, gate_mlp = temb.chunk(2, dim=1) + gate_msa, gate_mlp = gate_msa.unsqueeze(1), gate_mlp.unsqueeze(1) + return gate_msa, gate_mlp + + +class HunyuanImageCombinedTimeGuidanceEmbedding(nn.Module): + def __init__( + self, + embedding_dim: int, + guidance_embeds: bool = False, + use_meanflow: bool = False, + ): + super().__init__() + + self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + self.use_meanflow = use_meanflow + + self.time_proj_r = None + self.timestep_embedder_r = None + if use_meanflow: + self.time_proj_r = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder_r = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + self.guidance_embedder = None + if guidance_embeds: + self.guidance_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + def forward( + self, + timestep: torch.Tensor, + timestep_r: Optional[torch.Tensor] = None, + guidance: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=timestep.dtype)) + + if timestep_r is not None: + timesteps_proj_r = self.time_proj_r(timestep_r) + timesteps_emb_r = self.timestep_embedder_r(timesteps_proj_r.to(dtype=timestep.dtype)) + timesteps_emb = (timesteps_emb + timesteps_emb_r) / 2 + + if self.guidance_embedder is not None: + guidance_proj = self.time_proj(guidance) + guidance_emb = self.guidance_embedder(guidance_proj.to(dtype=timestep.dtype)) + conditioning = timesteps_emb + guidance_emb + else: + conditioning = timesteps_emb + + return conditioning + + +# IndividualTokenRefinerBlock +@maybe_allow_in_graph +class HunyuanImageIndividualTokenRefinerBlock(nn.Module): + def __init__( + self, + num_attention_heads: int, # 28 + attention_head_dim: int, # 128 + mlp_width_ratio: str = 4.0, + mlp_drop_rate: float = 0.0, + attention_bias: bool = True, + ) -> None: + super().__init__() + + hidden_size = num_attention_heads * attention_head_dim + + self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6) + self.attn = Attention( + query_dim=hidden_size, + cross_attention_dim=None, + heads=num_attention_heads, + dim_head=attention_head_dim, + bias=attention_bias, + ) + + self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6) + self.ff = FeedForward(hidden_size, mult=mlp_width_ratio, activation_fn="linear-silu", dropout=mlp_drop_rate) + + self.norm_out = HunyuanImageAdaNorm(hidden_size, 2 * hidden_size) + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + norm_hidden_states = self.norm1(hidden_states) + + attn_output = self.attn( + hidden_states=norm_hidden_states, + encoder_hidden_states=None, + attention_mask=attention_mask, + ) + + gate_msa, gate_mlp = self.norm_out(temb) + hidden_states = hidden_states + attn_output * gate_msa + + ff_output = self.ff(self.norm2(hidden_states)) + hidden_states = hidden_states + ff_output * gate_mlp + + return hidden_states + + +class HunyuanImageIndividualTokenRefiner(nn.Module): + def __init__( + self, + num_attention_heads: int, + attention_head_dim: int, + num_layers: int, + mlp_width_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + attention_bias: bool = True, + ) -> None: + super().__init__() + + self.refiner_blocks = nn.ModuleList( + [ + HunyuanImageIndividualTokenRefinerBlock( + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + mlp_width_ratio=mlp_width_ratio, + mlp_drop_rate=mlp_drop_rate, + attention_bias=attention_bias, + ) + for _ in range(num_layers) + ] + ) + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + ) -> None: + self_attn_mask = None + if attention_mask is not None: + batch_size = attention_mask.shape[0] + seq_len = attention_mask.shape[1] + attention_mask = attention_mask.to(hidden_states.device) + self_attn_mask_1 = attention_mask.view(batch_size, 1, 1, seq_len).repeat(1, 1, seq_len, 1) + self_attn_mask_2 = self_attn_mask_1.transpose(2, 3) + self_attn_mask = (self_attn_mask_1 & self_attn_mask_2).bool() + self_attn_mask[:, :, :, 0] = True + + for block in self.refiner_blocks: + hidden_states = block(hidden_states, temb, self_attn_mask) + + return hidden_states + + +# txt_in +class HunyuanImageTokenRefiner(nn.Module): + def __init__( + self, + in_channels: int, + num_attention_heads: int, + attention_head_dim: int, + num_layers: int, + mlp_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + attention_bias: bool = True, + ) -> None: + super().__init__() + + hidden_size = num_attention_heads * attention_head_dim + + self.time_text_embed = CombinedTimestepTextProjEmbeddings( + embedding_dim=hidden_size, pooled_projection_dim=in_channels + ) + self.proj_in = nn.Linear(in_channels, hidden_size, bias=True) + self.token_refiner = HunyuanImageIndividualTokenRefiner( + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + num_layers=num_layers, + mlp_width_ratio=mlp_ratio, + mlp_drop_rate=mlp_drop_rate, + attention_bias=attention_bias, + ) + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.LongTensor, + attention_mask: Optional[torch.LongTensor] = None, + ) -> torch.Tensor: + if attention_mask is None: + pooled_hidden_states = hidden_states.mean(dim=1) + else: + original_dtype = hidden_states.dtype + mask_float = attention_mask.float().unsqueeze(-1) + pooled_hidden_states = (hidden_states * mask_float).sum(dim=1) / mask_float.sum(dim=1) + pooled_hidden_states = pooled_hidden_states.to(original_dtype) + + temb = self.time_text_embed(timestep, pooled_hidden_states) + hidden_states = self.proj_in(hidden_states) + hidden_states = self.token_refiner(hidden_states, temb, attention_mask) + + return hidden_states + + +class HunyuanImageRotaryPosEmbed(nn.Module): + def __init__( + self, patch_size: Union[Tuple, List[int]], rope_dim: Union[Tuple, List[int]], theta: float = 256.0 + ) -> None: + super().__init__() + + if not isinstance(patch_size, (tuple, list)) or len(patch_size) not in [2, 3]: + raise ValueError(f"patch_size must be a tuple or list of length 2 or 3, got {patch_size}") + + if not isinstance(rope_dim, (tuple, list)) or len(rope_dim) not in [2, 3]: + raise ValueError(f"rope_dim must be a tuple or list of length 2 or 3, got {rope_dim}") + + if not len(patch_size) == len(rope_dim): + raise ValueError(f"patch_size and rope_dim must have the same length, got {patch_size} and {rope_dim}") + + self.patch_size = patch_size + self.rope_dim = rope_dim + self.theta = theta + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if hidden_states.ndim == 5: + _, _, frame, height, width = hidden_states.shape + patch_size_frame, patch_size_height, patch_size_width = self.patch_size + rope_sizes = [frame // patch_size_frame, height // patch_size_height, width // patch_size_width] + elif hidden_states.ndim == 4: + _, _, height, width = hidden_states.shape + patch_size_height, patch_size_width = self.patch_size + rope_sizes = [height // patch_size_height, width // patch_size_width] + else: + raise ValueError(f"hidden_states must be a 4D or 5D tensor, got {hidden_states.shape}") + + axes_grids = [] + for i in range(len(rope_sizes)): + grid = torch.arange(0, rope_sizes[i], device=hidden_states.device, dtype=torch.float32) + axes_grids.append(grid) + grid = torch.meshgrid(*axes_grids, indexing="ij") # dim x [H, W] + grid = torch.stack(grid, dim=0) # [2, H, W] + + freqs = [] + for i in range(len(rope_sizes)): + freq = get_1d_rotary_pos_embed(self.rope_dim[i], grid[i].reshape(-1), self.theta, use_real=True) + freqs.append(freq) + + freqs_cos = torch.cat([f[0] for f in freqs], dim=1) # (W * H * T, D / 2) + freqs_sin = torch.cat([f[1] for f in freqs], dim=1) # (W * H * T, D / 2) + return freqs_cos, freqs_sin + + +@maybe_allow_in_graph +class HunyuanImageSingleTransformerBlock(nn.Module): + def __init__( + self, + num_attention_heads: int, + attention_head_dim: int, + mlp_ratio: float = 4.0, + qk_norm: str = "rms_norm", + ) -> None: + super().__init__() + + hidden_size = num_attention_heads * attention_head_dim + mlp_dim = int(hidden_size * mlp_ratio) + + self.attn = Attention( + query_dim=hidden_size, + cross_attention_dim=None, + dim_head=attention_head_dim, + heads=num_attention_heads, + out_dim=hidden_size, + bias=True, + processor=HunyuanImageAttnProcessor(), + qk_norm=qk_norm, + eps=1e-6, + pre_only=True, + ) + + self.norm = AdaLayerNormZeroSingle(hidden_size, norm_type="layer_norm") + self.proj_mlp = nn.Linear(hidden_size, mlp_dim) + self.act_mlp = nn.GELU(approximate="tanh") + self.proj_out = nn.Linear(hidden_size + mlp_dim, hidden_size) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + *args, + **kwargs, + ) -> torch.Tensor: + text_seq_length = encoder_hidden_states.shape[1] + hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1) + + residual = hidden_states + + # 1. Input normalization + norm_hidden_states, gate = self.norm(hidden_states, emb=temb) + mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states)) + + norm_hidden_states, norm_encoder_hidden_states = ( + norm_hidden_states[:, :-text_seq_length, :], + norm_hidden_states[:, -text_seq_length:, :], + ) + + # 2. Attention + attn_output, context_attn_output = self.attn( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_encoder_hidden_states, + attention_mask=attention_mask, + image_rotary_emb=image_rotary_emb, + ) + attn_output = torch.cat([attn_output, context_attn_output], dim=1) + + # 3. Modulation and residual connection + hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2) + hidden_states = gate.unsqueeze(1) * self.proj_out(hidden_states) + hidden_states = hidden_states + residual + + hidden_states, encoder_hidden_states = ( + hidden_states[:, :-text_seq_length, :], + hidden_states[:, -text_seq_length:, :], + ) + return hidden_states, encoder_hidden_states + + +@maybe_allow_in_graph +class HunyuanImageTransformerBlock(nn.Module): + def __init__( + self, + num_attention_heads: int, + attention_head_dim: int, + mlp_ratio: float, + qk_norm: str = "rms_norm", + ) -> None: + super().__init__() + + hidden_size = num_attention_heads * attention_head_dim + + self.norm1 = AdaLayerNormZero(hidden_size, norm_type="layer_norm") + self.norm1_context = AdaLayerNormZero(hidden_size, norm_type="layer_norm") + + self.attn = Attention( + query_dim=hidden_size, + cross_attention_dim=None, + added_kv_proj_dim=hidden_size, + dim_head=attention_head_dim, + heads=num_attention_heads, + out_dim=hidden_size, + context_pre_only=False, + bias=True, + processor=HunyuanImageAttnProcessor(), + qk_norm=qk_norm, + eps=1e-6, + ) + + self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.ff = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu-approximate") + + self.norm2_context = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.ff_context = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu-approximate") + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + *args, + **kwargs, + ) -> Tuple[torch.Tensor, torch.Tensor]: + # 1. Input normalization + norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb) + norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context( + encoder_hidden_states, emb=temb + ) + + # 2. Joint attention + attn_output, context_attn_output = self.attn( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_encoder_hidden_states, + attention_mask=attention_mask, + image_rotary_emb=image_rotary_emb, + ) + + # 3. Modulation and residual connection + hidden_states = hidden_states + attn_output * gate_msa.unsqueeze(1) + encoder_hidden_states = encoder_hidden_states + context_attn_output * c_gate_msa.unsqueeze(1) + + norm_hidden_states = self.norm2(hidden_states) + norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) + + norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] + + # 4. Feed-forward + ff_output = self.ff(norm_hidden_states) + context_ff_output = self.ff_context(norm_encoder_hidden_states) + + hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff_output + encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output + + return hidden_states, encoder_hidden_states + + +class HunyuanImageTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, CacheMixin): + r""" + The Transformer model used in [HunyuanImage-2.1](https://github.com/Tencent-Hunyuan/HunyuanImage-2.1). + + Args: + in_channels (`int`, defaults to `16`): + The number of channels in the input. + out_channels (`int`, defaults to `16`): + The number of channels in the output. + num_attention_heads (`int`, defaults to `24`): + The number of heads to use for multi-head attention. + attention_head_dim (`int`, defaults to `128`): + The number of channels in each head. + num_layers (`int`, defaults to `20`): + The number of layers of dual-stream blocks to use. + num_single_layers (`int`, defaults to `40`): + The number of layers of single-stream blocks to use. + num_refiner_layers (`int`, defaults to `2`): + The number of layers of refiner blocks to use. + mlp_ratio (`float`, defaults to `4.0`): + The ratio of the hidden layer size to the input size in the feedforward network. + patch_size (`int`, defaults to `2`): + The size of the spatial patches to use in the patch embedding layer. + patch_size_t (`int`, defaults to `1`): + The size of the tmeporal patches to use in the patch embedding layer. + qk_norm (`str`, defaults to `rms_norm`): + The normalization to use for the query and key projections in the attention layers. + guidance_embeds (`bool`, defaults to `True`): + Whether to use guidance embeddings in the model. + text_embed_dim (`int`, defaults to `4096`): + Input dimension of text embeddings from the text encoder. + pooled_projection_dim (`int`, defaults to `768`): + The dimension of the pooled projection of the text embeddings. + rope_theta (`float`, defaults to `256.0`): + The value of theta to use in the RoPE layer. + rope_axes_dim (`Tuple[int]`, defaults to `(16, 56, 56)`): + The dimensions of the axes to use in the RoPE layer. + image_condition_type (`str`, *optional*, defaults to `None`): + The type of image conditioning to use. If `None`, no image conditioning is used. If `latent_concat`, the + image is concatenated to the latent stream. If `token_replace`, the image is used to replace first-frame + tokens in the latent stream and apply conditioning. + """ + + _supports_gradient_checkpointing = True + _skip_layerwise_casting_patterns = ["x_embedder", "context_embedder", "norm"] + _no_split_modules = [ + "HunyuanImageTransformerBlock", + "HunyuanImageSingleTransformerBlock", + "HunyuanImagePatchEmbed", + "HunyuanImageTokenRefiner", + ] + _repeated_blocks = [ + "HunyuanImageTransformerBlock", + "HunyuanImageSingleTransformerBlock", + ] + + @register_to_config + def __init__( + self, + in_channels: int = 64, + out_channels: int = 64, + num_attention_heads: int = 28, + attention_head_dim: int = 128, + num_layers: int = 20, + num_single_layers: int = 40, + num_refiner_layers: int = 2, + mlp_ratio: float = 4.0, + patch_size: Tuple[int, int] = (1, 1), + qk_norm: str = "rms_norm", + guidance_embeds: bool = False, + text_embed_dim: int = 3584, + text_embed_2_dim: Optional[int] = None, + rope_theta: float = 256.0, + rope_axes_dim: Tuple[int] = (64, 64), + use_meanflow: bool = False, + ) -> None: + super().__init__() + + if not (isinstance(patch_size, (tuple, list)) and len(patch_size) in [2, 3]): + raise ValueError(f"patch_size must be a tuple of length 2 or 3, got {patch_size}") + + inner_dim = num_attention_heads * attention_head_dim + out_channels = out_channels or in_channels + + # 1. Latent and condition embedders + self.x_embedder = HunyuanImagePatchEmbed(patch_size, in_channels, inner_dim) + self.context_embedder = HunyuanImageTokenRefiner( + text_embed_dim, num_attention_heads, attention_head_dim, num_layers=num_refiner_layers + ) + + if text_embed_2_dim is not None: + self.context_embedder_2 = HunyuanImageByT5TextProjection(text_embed_2_dim, 2048, inner_dim) + else: + self.context_embedder_2 = None + + self.time_guidance_embed = HunyuanImageCombinedTimeGuidanceEmbedding(inner_dim, guidance_embeds, use_meanflow) + + # 2. RoPE + self.rope = HunyuanImageRotaryPosEmbed(patch_size, rope_axes_dim, rope_theta) + + # 3. Dual stream transformer blocks + + self.transformer_blocks = nn.ModuleList( + [ + HunyuanImageTransformerBlock( + num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm + ) + for _ in range(num_layers) + ] + ) + + # 4. Single stream transformer blocks + self.single_transformer_blocks = nn.ModuleList( + [ + HunyuanImageSingleTransformerBlock( + num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm + ) + for _ in range(num_single_layers) + ] + ) + + # 5. Output projection + self.norm_out = AdaLayerNormContinuous(inner_dim, inner_dim, elementwise_affine=False, eps=1e-6) + self.proj_out = nn.Linear(inner_dim, math.prod(patch_size) * out_channels) + + self.gradient_checkpointing = False + + @property + # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors + def attn_processors(self) -> Dict[str, AttentionProcessor]: + r""" + Returns: + `dict` of attention processors: A dictionary containing all attention processors used in the model with + indexed by its weight name. + """ + # set recursively + processors = {} + + def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]): + if hasattr(module, "get_processor"): + processors[f"{name}.processor"] = module.get_processor() + + for sub_name, child in module.named_children(): + fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) + + return processors + + for name, module in self.named_children(): + fn_recursive_add_processors(name, module, processors) + + return processors + + # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor + def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): + r""" + Sets the attention processor to use to compute attention. + + Parameters: + processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): + The instantiated processor class or a dictionary of processor classes that will be set as the processor + for **all** `Attention` layers. + + If `processor` is a dict, the key needs to define the path to the corresponding cross attention + processor. This is strongly recommended when setting trainable attention processors. + + """ + count = len(self.attn_processors.keys()) + + if isinstance(processor, dict) and len(processor) != count: + raise ValueError( + f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" + f" number of attention layers: {count}. Please make sure to pass {count} processor classes." + ) + + def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): + if hasattr(module, "set_processor"): + if not isinstance(processor, dict): + module.set_processor(processor) + else: + module.set_processor(processor.pop(f"{name}.processor")) + + for sub_name, child in module.named_children(): + fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) + + for name, module in self.named_children(): + fn_recursive_attn_processor(name, module, processor) + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.LongTensor, + encoder_hidden_states: torch.Tensor, + encoder_attention_mask: torch.Tensor, + timestep_r: Optional[torch.LongTensor] = None, + encoder_hidden_states_2: Optional[torch.Tensor] = None, + encoder_attention_mask_2: Optional[torch.Tensor] = None, + guidance: Optional[torch.Tensor] = None, + attention_kwargs: Optional[Dict[str, Any]] = None, + return_dict: bool = True, + ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: + if attention_kwargs is not None: + attention_kwargs = attention_kwargs.copy() + lora_scale = attention_kwargs.pop("scale", 1.0) + else: + lora_scale = 1.0 + + if USE_PEFT_BACKEND: + # weight the lora layers by setting `lora_scale` for each PEFT layer + scale_lora_layers(self, lora_scale) + else: + if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: + logger.warning( + "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective." + ) + + if hidden_states.ndim == 4: + batch_size, channels, height, width = hidden_states.shape + sizes = (height, width) + elif hidden_states.ndim == 5: + batch_size, channels, frame, height, width = hidden_states.shape + sizes = (frame, height, width) + else: + raise ValueError(f"hidden_states must be a 4D or 5D tensor, got {hidden_states.shape}") + + post_patch_sizes = tuple(d // p for d, p in zip(sizes, self.config.patch_size)) + + # 1. RoPE + image_rotary_emb = self.rope(hidden_states) + + # 2. Conditional embeddings + encoder_attention_mask = encoder_attention_mask.bool() + temb = self.time_guidance_embed(timestep, guidance=guidance, timestep_r=timestep_r) + hidden_states = self.x_embedder(hidden_states) + encoder_hidden_states = self.context_embedder(encoder_hidden_states, timestep, encoder_attention_mask) + + if self.context_embedder_2 is not None and encoder_hidden_states_2 is not None: + encoder_hidden_states_2 = self.context_embedder_2(encoder_hidden_states_2) + + encoder_attention_mask_2 = encoder_attention_mask_2.bool() + + # reorder and combine text tokens: combine valid tokens first, then padding + new_encoder_hidden_states = [] + new_encoder_attention_mask = [] + + for text, text_mask, text_2, text_mask_2 in zip( + encoder_hidden_states, encoder_attention_mask, encoder_hidden_states_2, encoder_attention_mask_2 + ): + # Concatenate: [valid_mllm, valid_byt5, invalid_mllm, invalid_byt5] + new_encoder_hidden_states.append( + torch.cat( + [ + text_2[text_mask_2], # valid byt5 + text[text_mask], # valid mllm + text_2[~text_mask_2], # invalid byt5 + text[~text_mask], # invalid mllm + ], + dim=0, + ) + ) + + # Apply same reordering to attention masks + new_encoder_attention_mask.append( + torch.cat( + [ + text_mask_2[text_mask_2], + text_mask[text_mask], + text_mask_2[~text_mask_2], + text_mask[~text_mask], + ], + dim=0, + ) + ) + + encoder_hidden_states = torch.stack(new_encoder_hidden_states) + encoder_attention_mask = torch.stack(new_encoder_attention_mask) + + attention_mask = torch.nn.functional.pad(encoder_attention_mask, (hidden_states.shape[1], 0), value=True) + attention_mask = attention_mask.unsqueeze(1).unsqueeze(2) + # 3. Transformer blocks + if torch.is_grad_enabled() and self.gradient_checkpointing: + for block in self.transformer_blocks: + hidden_states, encoder_hidden_states = self._gradient_checkpointing_func( + block, + hidden_states, + encoder_hidden_states, + temb, + attention_mask=attention_mask, + image_rotary_emb=image_rotary_emb, + ) + + for block in self.single_transformer_blocks: + hidden_states, encoder_hidden_states = self._gradient_checkpointing_func( + block, + hidden_states, + encoder_hidden_states, + temb, + attention_mask=attention_mask, + image_rotary_emb=image_rotary_emb, + ) + + else: + for block in self.transformer_blocks: + hidden_states, encoder_hidden_states = block( + hidden_states, + encoder_hidden_states, + temb, + attention_mask=attention_mask, + image_rotary_emb=image_rotary_emb, + ) + + for block in self.single_transformer_blocks: + hidden_states, encoder_hidden_states = block( + hidden_states, + encoder_hidden_states, + temb, + attention_mask=attention_mask, + image_rotary_emb=image_rotary_emb, + ) + + # 4. Output projection + hidden_states = self.norm_out(hidden_states, temb) + hidden_states = self.proj_out(hidden_states) + + # 5. unpatchify + # reshape: [batch_size, *post_patch_dims, channels, *patch_size] + out_channels = self.config.out_channels + reshape_dims = [batch_size] + list(post_patch_sizes) + [out_channels] + list(self.config.patch_size) + hidden_states = hidden_states.reshape(*reshape_dims) + + # create permutation pattern: batch, channels, then interleave post_patch and patch dims + # For 4D: [0, 3, 1, 4, 2, 5] -> batch, channels, post_patch_height, patch_size_height, post_patch_width, patch_size_width + # For 5D: [0, 4, 1, 5, 2, 6, 3, 7] -> batch, channels, post_patch_frame, patch_size_frame, post_patch_height, patch_size_height, post_patch_width, patch_size_width + ndim = len(post_patch_sizes) + permute_pattern = [0, ndim + 1] # batch, channels + for i in range(ndim): + permute_pattern.extend([i + 1, ndim + 2 + i]) # post_patch_sizes[i], patch_sizes[i] + hidden_states = hidden_states.permute(*permute_pattern) + + # flatten patch dimensions: flatten each (post_patch_size, patch_size) pair + # batch_size, channels, post_patch_sizes[0] * patch_sizes[0], post_patch_sizes[1] * patch_sizes[1], ... + final_dims = [batch_size, out_channels] + [ + post_patch * patch for post_patch, patch in zip(post_patch_sizes, self.config.patch_size) + ] + hidden_states = hidden_states.reshape(*final_dims) + + if USE_PEFT_BACKEND: + # remove `lora_scale` from each PEFT layer + unscale_lora_layers(self, lora_scale) + + if not return_dict: + return (hidden_states,) + + return Transformer2DModelOutput(sample=hidden_states) diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index cfbca48a98..bf067555a8 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -130,8 +130,14 @@ class PipelineState: Allow attribute access to intermediate values. If an attribute is not found in the object, look for it in the intermediates dict. """ - if name in self.values: - return self.values[name] + # Use object.__getattribute__ to avoid infinite recursion during deepcopy + try: + values = object.__getattribute__(self, "values") + except AttributeError: + raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'") + + if name in values: + return values[name] raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'") def __repr__(self): @@ -2492,6 +2498,8 @@ class ModularPipeline(ConfigMixin, PushToHubMixin): """ if state is None: state = PipelineState() + else: + state = deepcopy(state) # Make a copy of the input kwargs passed_kwargs = kwargs.copy() diff --git a/src/diffusers/modular_pipelines/qwenimage/denoise.py b/src/diffusers/modular_pipelines/qwenimage/denoise.py index d0704ee6e0..49acd2dc02 100644 --- a/src/diffusers/modular_pipelines/qwenimage/denoise.py +++ b/src/diffusers/modular_pipelines/qwenimage/denoise.py @@ -238,19 +238,27 @@ class QwenImageLoopDenoiser(ModularPipelineBlocks): @torch.no_grad() def __call__(self, components: QwenImageModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): - guider_input_fields = { - "encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds"), - "encoder_hidden_states_mask": ("prompt_embeds_mask", "negative_prompt_embeds_mask"), - "txt_seq_lens": ("txt_seq_lens", "negative_txt_seq_lens"), + guider_inputs = { + "encoder_hidden_states": ( + getattr(block_state, "prompt_embeds", None), + getattr(block_state, "negative_prompt_embeds", None), + ), + "encoder_hidden_states_mask": ( + getattr(block_state, "prompt_embeds_mask", None), + getattr(block_state, "negative_prompt_embeds_mask", None), + ), + "txt_seq_lens": ( + getattr(block_state, "txt_seq_lens", None), + getattr(block_state, "negative_txt_seq_lens", None), + ), } components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) - guider_state = components.guider.prepare_inputs(block_state, guider_input_fields) + guider_state = components.guider.prepare_inputs(guider_inputs) for guider_state_batch in guider_state: components.guider.prepare_models(components.transformer) - cond_kwargs = guider_state_batch.as_dict() - cond_kwargs = {k: v for k, v in cond_kwargs.items() if k in guider_input_fields} + cond_kwargs = {input_name: getattr(guider_state_batch, input_name) for input_name in guider_inputs.keys()} # YiYi TODO: add cache context guider_state_batch.noise_pred = components.transformer( @@ -328,19 +336,27 @@ class QwenImageEditLoopDenoiser(ModularPipelineBlocks): @torch.no_grad() def __call__(self, components: QwenImageModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): - guider_input_fields = { - "encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds"), - "encoder_hidden_states_mask": ("prompt_embeds_mask", "negative_prompt_embeds_mask"), - "txt_seq_lens": ("txt_seq_lens", "negative_txt_seq_lens"), + guider_inputs = { + "encoder_hidden_states": ( + getattr(block_state, "prompt_embeds", None), + getattr(block_state, "negative_prompt_embeds", None), + ), + "encoder_hidden_states_mask": ( + getattr(block_state, "prompt_embeds_mask", None), + getattr(block_state, "negative_prompt_embeds_mask", None), + ), + "txt_seq_lens": ( + getattr(block_state, "txt_seq_lens", None), + getattr(block_state, "negative_txt_seq_lens", None), + ), } components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) - guider_state = components.guider.prepare_inputs(block_state, guider_input_fields) + guider_state = components.guider.prepare_inputs(guider_inputs) for guider_state_batch in guider_state: components.guider.prepare_models(components.transformer) - cond_kwargs = guider_state_batch.as_dict() - cond_kwargs = {k: v for k, v in cond_kwargs.items() if k in guider_input_fields} + cond_kwargs = {input_name: getattr(guider_state_batch, input_name) for input_name in guider_inputs.keys()} # YiYi TODO: add cache context guider_state_batch.noise_pred = components.transformer( diff --git a/src/diffusers/modular_pipelines/stable_diffusion_xl/denoise.py b/src/diffusers/modular_pipelines/stable_diffusion_xl/denoise.py index 83e6007e3f..862315e591 100644 --- a/src/diffusers/modular_pipelines/stable_diffusion_xl/denoise.py +++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/denoise.py @@ -201,27 +201,41 @@ class StableDiffusionXLLoopDenoiser(ModularPipelineBlocks): ) -> PipelineState: # Map the keys we'll see on each `guider_state_batch` (e.g. guider_state_batch.prompt_embeds) # to the corresponding (cond, uncond) fields on block_state. (e.g. block_state.prompt_embeds, block_state.negative_prompt_embeds) - guider_input_fields = { - "prompt_embeds": ("prompt_embeds", "negative_prompt_embeds"), - "time_ids": ("add_time_ids", "negative_add_time_ids"), - "text_embeds": ("pooled_prompt_embeds", "negative_pooled_prompt_embeds"), - "image_embeds": ("ip_adapter_embeds", "negative_ip_adapter_embeds"), + guider_inputs = { + "prompt_embeds": ( + getattr(block_state, "prompt_embeds", None), + getattr(block_state, "negative_prompt_embeds", None), + ), + "time_ids": ( + getattr(block_state, "add_time_ids", None), + getattr(block_state, "negative_add_time_ids", None), + ), + "text_embeds": ( + getattr(block_state, "pooled_prompt_embeds", None), + getattr(block_state, "negative_pooled_prompt_embeds", None), + ), + "image_embeds": ( + getattr(block_state, "ip_adapter_embeds", None), + getattr(block_state, "negative_ip_adapter_embeds", None), + ), } components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) - # Prepare mini‐batches according to guidance method and `guider_input_fields` - # Each guider_state_batch will have .prompt_embeds, .time_ids, text_embeds, image_embeds. - # e.g. for CFG, we prepare two batches: one for uncond, one for cond - # for first batch, guider_state_batch.prompt_embeds correspond to block_state.prompt_embeds - # for second batch, guider_state_batch.prompt_embeds correspond to block_state.negative_prompt_embeds - guider_state = components.guider.prepare_inputs(block_state, guider_input_fields) + # The guider splits model inputs into separate batches for conditional/unconditional predictions. + # For CFG with guider_inputs = {"encoder_hidden_states": (prompt_embeds, negative_prompt_embeds)}: + # you will get a guider_state with two batches: + # guider_state = [ + # {"encoder_hidden_states": prompt_embeds, "__guidance_identifier__": "pred_cond"}, # conditional batch + # {"encoder_hidden_states": negative_prompt_embeds, "__guidance_identifier__": "pred_uncond"}, # unconditional batch + # ] + # Other guidance methods may return 1 batch (no guidance) or 3+ batches (e.g., PAG, APG). + guider_state = components.guider.prepare_inputs(guider_inputs) # run the denoiser for each guidance batch for guider_state_batch in guider_state: components.guider.prepare_models(components.unet) - cond_kwargs = guider_state_batch.as_dict() - cond_kwargs = {k: v for k, v in cond_kwargs.items() if k in guider_input_fields} + cond_kwargs = {input_name: getattr(guider_state_batch, input_name) for input_name in guider_inputs.keys()} prompt_embeds = cond_kwargs.pop("prompt_embeds") # Predict the noise residual @@ -344,11 +358,23 @@ class StableDiffusionXLControlNetLoopDenoiser(ModularPipelineBlocks): # Map the keys we'll see on each `guider_state_batch` (e.g. guider_state_batch.prompt_embeds) # to the corresponding (cond, uncond) fields on block_state. (e.g. block_state.prompt_embeds, block_state.negative_prompt_embeds) - guider_input_fields = { - "prompt_embeds": ("prompt_embeds", "negative_prompt_embeds"), - "time_ids": ("add_time_ids", "negative_add_time_ids"), - "text_embeds": ("pooled_prompt_embeds", "negative_pooled_prompt_embeds"), - "image_embeds": ("ip_adapter_embeds", "negative_ip_adapter_embeds"), + guider_inputs = { + "prompt_embeds": ( + getattr(block_state, "prompt_embeds", None), + getattr(block_state, "negative_prompt_embeds", None), + ), + "time_ids": ( + getattr(block_state, "add_time_ids", None), + getattr(block_state, "negative_add_time_ids", None), + ), + "text_embeds": ( + getattr(block_state, "pooled_prompt_embeds", None), + getattr(block_state, "negative_pooled_prompt_embeds", None), + ), + "image_embeds": ( + getattr(block_state, "ip_adapter_embeds", None), + getattr(block_state, "negative_ip_adapter_embeds", None), + ), } # cond_scale for the timestep (controlnet input) @@ -369,12 +395,15 @@ class StableDiffusionXLControlNetLoopDenoiser(ModularPipelineBlocks): # guided denoiser step components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) - # Prepare mini‐batches according to guidance method and `guider_input_fields` - # Each guider_state_batch will have .prompt_embeds, .time_ids, text_embeds, image_embeds. - # e.g. for CFG, we prepare two batches: one for uncond, one for cond - # for first batch, guider_state_batch.prompt_embeds correspond to block_state.prompt_embeds - # for second batch, guider_state_batch.prompt_embeds correspond to block_state.negative_prompt_embeds - guider_state = components.guider.prepare_inputs(block_state, guider_input_fields) + # The guider splits model inputs into separate batches for conditional/unconditional predictions. + # For CFG with guider_inputs = {"encoder_hidden_states": (prompt_embeds, negative_prompt_embeds)}: + # you will get a guider_state with two batches: + # guider_state = [ + # {"encoder_hidden_states": prompt_embeds, "__guidance_identifier__": "pred_cond"}, # conditional batch + # {"encoder_hidden_states": negative_prompt_embeds, "__guidance_identifier__": "pred_uncond"}, # unconditional batch + # ] + # Other guidance methods may return 1 batch (no guidance) or 3+ batches (e.g., PAG, APG). + guider_state = components.guider.prepare_inputs(guider_inputs) # run the denoiser for each guidance batch for guider_state_batch in guider_state: diff --git a/src/diffusers/modular_pipelines/wan/denoise.py b/src/diffusers/modular_pipelines/wan/denoise.py index 66c51493bd..4f3ca80acc 100644 --- a/src/diffusers/modular_pipelines/wan/denoise.py +++ b/src/diffusers/modular_pipelines/wan/denoise.py @@ -94,25 +94,30 @@ class WanLoopDenoiser(ModularPipelineBlocks): ) -> PipelineState: # Map the keys we'll see on each `guider_state_batch` (e.g. guider_state_batch.prompt_embeds) # to the corresponding (cond, uncond) fields on block_state. (e.g. block_state.prompt_embeds, block_state.negative_prompt_embeds) - guider_input_fields = { - "prompt_embeds": ("prompt_embeds", "negative_prompt_embeds"), + guider_inputs = { + "prompt_embeds": ( + getattr(block_state, "prompt_embeds", None), + getattr(block_state, "negative_prompt_embeds", None), + ), } transformer_dtype = components.transformer.dtype components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) - # Prepare mini‐batches according to guidance method and `guider_input_fields` - # Each guider_state_batch will have .prompt_embeds, .time_ids, text_embeds, image_embeds. - # e.g. for CFG, we prepare two batches: one for uncond, one for cond - # for first batch, guider_state_batch.prompt_embeds correspond to block_state.prompt_embeds - # for second batch, guider_state_batch.prompt_embeds correspond to block_state.negative_prompt_embeds - guider_state = components.guider.prepare_inputs(block_state, guider_input_fields) + # The guider splits model inputs into separate batches for conditional/unconditional predictions. + # For CFG with guider_inputs = {"encoder_hidden_states": (prompt_embeds, negative_prompt_embeds)}: + # you will get a guider_state with two batches: + # guider_state = [ + # {"encoder_hidden_states": prompt_embeds, "__guidance_identifier__": "pred_cond"}, # conditional batch + # {"encoder_hidden_states": negative_prompt_embeds, "__guidance_identifier__": "pred_uncond"}, # unconditional batch + # ] + # Other guidance methods may return 1 batch (no guidance) or 3+ batches (e.g., PAG, APG). + guider_state = components.guider.prepare_inputs(guider_inputs) # run the denoiser for each guidance batch for guider_state_batch in guider_state: components.guider.prepare_models(components.transformer) - cond_kwargs = guider_state_batch.as_dict() - cond_kwargs = {k: v for k, v in cond_kwargs.items() if k in guider_input_fields} + cond_kwargs = {input_name: getattr(guider_state_batch, input_name) for input_name in guider_inputs.keys()} prompt_embeds = cond_kwargs.pop("prompt_embeds") # Predict the noise residual diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index ff64958d46..65e8b2469e 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -241,6 +241,7 @@ else: "HunyuanVideoImageToVideoPipeline", "HunyuanVideoFramepackPipeline", ] + _import_structure["hunyuan_image"] = ["HunyuanImagePipeline", "HunyuanImageRefinerPipeline"] _import_structure["kandinsky"] = [ "KandinskyCombinedPipeline", "KandinskyImg2ImgCombinedPipeline", @@ -640,6 +641,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: ReduxImageEncoder, ) from .hidream_image import HiDreamImagePipeline + from .hunyuan_image import HunyuanImagePipeline, HunyuanImageRefinerPipeline from .hunyuan_video import ( HunyuanSkyreelsImageToVideoPipeline, HunyuanVideoFramepackPipeline, diff --git a/src/diffusers/pipelines/hunyuan_image/__init__.py b/src/diffusers/pipelines/hunyuan_image/__init__.py new file mode 100644 index 0000000000..7da72fa12b --- /dev/null +++ b/src/diffusers/pipelines/hunyuan_image/__init__.py @@ -0,0 +1,50 @@ +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects # noqa F403 + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["pipeline_hunyuanimage"] = ["HunyuanImagePipeline"] + _import_structure["pipeline_hunyuanimage_refiner"] = ["HunyuanImageRefinerPipeline"] + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() + + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * + else: + from .pipeline_hunyuanimage import HunyuanImagePipeline + from .pipeline_hunyuanimage_refiner import HunyuanImageRefinerPipeline + +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) + + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/hunyuan_image/pipeline_hunyuanimage.py b/src/diffusers/pipelines/hunyuan_image/pipeline_hunyuanimage.py new file mode 100644 index 0000000000..658935ccd8 --- /dev/null +++ b/src/diffusers/pipelines/hunyuan_image/pipeline_hunyuanimage.py @@ -0,0 +1,866 @@ +# Copyright 2025 Hunyuan-Image Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import re +from typing import Any, Callable, Dict, List, Optional, Union + +import numpy as np +import torch +from transformers import ByT5Tokenizer, Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, T5EncoderModel + +from ...guiders import AdaptiveProjectedMixGuidance +from ...image_processor import VaeImageProcessor +from ...models import AutoencoderKLHunyuanImage, HunyuanImageTransformer2DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import is_torch_xla_available, logging, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ..pipeline_utils import DiffusionPipeline +from .pipeline_output import HunyuanImagePipelineOutput + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import HunyuanImagePipeline + + >>> pipe = HunyuanImagePipeline.from_pretrained( + ... "hunyuanvideo-community/HunyuanImage-2.1-Diffusers", torch_dtype=torch.bfloat16 + ... ) + >>> pipe.to("cuda") + >>> prompt = "A cat holding a sign that says hello world" + >>> # Depending on the variant being used, the pipeline call will slightly vary. + >>> # Refer to the pipeline documentation for more details. + >>> image = pipe(prompt, negative_prompt="", num_inference_steps=50).images[0] + >>> image.save("hunyuanimage.png") + ``` +""" + + +def extract_glyph_text(prompt: str): + """ + Extract text enclosed in quotes for glyph rendering. + + Finds text in single quotes, double quotes, and Chinese quotes, then formats it for byT5 processing. + + Args: + prompt: Input text prompt + + Returns: + Formatted glyph text string or None if no quoted text found + """ + text_prompt_texts = [] + pattern_quote_single = r"\'(.*?)\'" + pattern_quote_double = r"\"(.*?)\"" + pattern_quote_chinese_single = r"‘(.*?)’" + pattern_quote_chinese_double = r"“(.*?)”" + + matches_quote_single = re.findall(pattern_quote_single, prompt) + matches_quote_double = re.findall(pattern_quote_double, prompt) + matches_quote_chinese_single = re.findall(pattern_quote_chinese_single, prompt) + matches_quote_chinese_double = re.findall(pattern_quote_chinese_double, prompt) + + text_prompt_texts.extend(matches_quote_single) + text_prompt_texts.extend(matches_quote_double) + text_prompt_texts.extend(matches_quote_chinese_single) + text_prompt_texts.extend(matches_quote_chinese_double) + + if text_prompt_texts: + glyph_text_formatted = ". ".join([f'Text "{text}"' for text in text_prompt_texts]) + ". " + else: + glyph_text_formatted = None + + return glyph_text_formatted + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + timesteps: Optional[List[int]] = None, + sigmas: Optional[List[float]] = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`List[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`List[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class HunyuanImagePipeline(DiffusionPipeline): + r""" + The HunyuanImage pipeline for text-to-image generation. + + Args: + transformer ([`HunyuanImageTransformer2DModel`]): + Conditional Transformer (MMDiT) architecture to denoise the encoded image latents. + scheduler ([`FlowMatchEulerDiscreteScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded image latents. + vae ([`AutoencoderKLHunyuanImage`]): + Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. + text_encoder ([`Qwen2.5-VL-7B-Instruct`]): + [Qwen2.5-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct), specifically the + [Qwen2.5-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct) variant. + tokenizer (`Qwen2Tokenizer`): Tokenizer of class [Qwen2Tokenizer]. + text_encoder_2 ([`T5EncoderModel`]): + [T5EncoderModel](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel) + variant. + tokenizer_2 (`ByT5Tokenizer`): Tokenizer of class [ByT5Tokenizer] + guider ([`AdaptiveProjectedMixGuidance`]): + [AdaptiveProjectedMixGuidance]to be used to guide the image generation. + ocr_guider ([`AdaptiveProjectedMixGuidance`], *optional*): + [AdaptiveProjectedMixGuidance] to be used to guide the image generation when text rendering is needed. + """ + + model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae" + _callback_tensor_inputs = ["latents", "prompt_embeds"] + _optional_components = ["ocr_guider", "guider"] + + def __init__( + self, + scheduler: FlowMatchEulerDiscreteScheduler, + vae: AutoencoderKLHunyuanImage, + text_encoder: Qwen2_5_VLForConditionalGeneration, + tokenizer: Qwen2Tokenizer, + text_encoder_2: T5EncoderModel, + tokenizer_2: ByT5Tokenizer, + transformer: HunyuanImageTransformer2DModel, + guider: Optional[AdaptiveProjectedMixGuidance] = None, + ocr_guider: Optional[AdaptiveProjectedMixGuidance] = None, + ): + super().__init__() + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + text_encoder_2=text_encoder_2, + tokenizer_2=tokenizer_2, + transformer=transformer, + scheduler=scheduler, + guider=guider, + ocr_guider=ocr_guider, + ) + + self.vae_scale_factor = self.vae.config.spatial_compression_ratio if getattr(self, "vae", None) else 32 + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + self.tokenizer_max_length = 1000 + self.tokenizer_2_max_length = 128 + self.prompt_template_encode = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>" + self.prompt_template_encode_start_idx = 34 + self.default_sample_size = 64 + + def _get_qwen_prompt_embeds( + self, + tokenizer: Qwen2Tokenizer, + text_encoder: Qwen2_5_VLForConditionalGeneration, + prompt: Union[str, List[str]] = None, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + tokenizer_max_length: int = 1000, + template: str = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>", + drop_idx: int = 34, + hidden_state_skip_layer: int = 2, + ): + device = device or self._execution_device + dtype = dtype or text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + + txt = [template.format(e) for e in prompt] + txt_tokens = tokenizer( + txt, max_length=tokenizer_max_length + drop_idx, padding="max_length", truncation=True, return_tensors="pt" + ).to(device) + + encoder_hidden_states = text_encoder( + input_ids=txt_tokens.input_ids, + attention_mask=txt_tokens.attention_mask, + output_hidden_states=True, + ) + prompt_embeds = encoder_hidden_states.hidden_states[-(hidden_state_skip_layer + 1)] + + prompt_embeds = prompt_embeds[:, drop_idx:] + encoder_attention_mask = txt_tokens.attention_mask[:, drop_idx:] + + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + encoder_attention_mask = encoder_attention_mask.to(device=device) + + return prompt_embeds, encoder_attention_mask + + def _get_byt5_prompt_embeds( + self, + tokenizer: ByT5Tokenizer, + text_encoder: T5EncoderModel, + prompt: str, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + tokenizer_max_length: int = 128, + ): + device = device or self._execution_device + dtype = dtype or text_encoder.dtype + + if isinstance(prompt, list): + raise ValueError("byt5 prompt should be a string") + elif prompt is None: + raise ValueError("byt5 prompt should not be None") + + txt_tokens = tokenizer( + prompt, + padding="max_length", + max_length=tokenizer_max_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ).to(device) + + prompt_embeds = text_encoder( + input_ids=txt_tokens.input_ids, + attention_mask=txt_tokens.attention_mask.float(), + )[0] + + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + encoder_attention_mask = txt_tokens.attention_mask.to(device=device) + + return prompt_embeds, encoder_attention_mask + + def encode_prompt( + self, + prompt: Union[str, List[str]], + device: Optional[torch.device] = None, + batch_size: int = 1, + num_images_per_prompt: int = 1, + prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds_mask: Optional[torch.Tensor] = None, + prompt_embeds_2: Optional[torch.Tensor] = None, + prompt_embeds_mask_2: Optional[torch.Tensor] = None, + ): + r""" + + Args: + prompt (`str` or `List[str]`, *optional*): + prompt to be encoded + device: (`torch.device`): + torch device + batch_size (`int`): + batch size of prompts, defaults to 1 + num_images_per_prompt (`int`): + number of images that should be generated per prompt + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. If not provided, text embeddings will be generated from `prompt` input + argument. + prompt_embeds_mask (`torch.Tensor`, *optional*): + Pre-generated text mask. If not provided, text mask will be generated from `prompt` input argument. + prompt_embeds_2 (`torch.Tensor`, *optional*): + Pre-generated glyph text embeddings from ByT5. If not provided, will be generated from `prompt` input + argument using self.tokenizer_2 and self.text_encoder_2. + prompt_embeds_mask_2 (`torch.Tensor`, *optional*): + Pre-generated glyph text mask from ByT5. If not provided, will be generated from `prompt` input + argument using self.tokenizer_2 and self.text_encoder_2. + """ + device = device or self._execution_device + + if prompt is None: + prompt = [""] * batch_size + + prompt = [prompt] if isinstance(prompt, str) else prompt + + if prompt_embeds is None: + prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds( + tokenizer=self.tokenizer, + text_encoder=self.text_encoder, + prompt=prompt, + device=device, + tokenizer_max_length=self.tokenizer_max_length, + template=self.prompt_template_encode, + drop_idx=self.prompt_template_encode_start_idx, + ) + + if prompt_embeds_2 is None: + prompt_embeds_2_list = [] + prompt_embeds_mask_2_list = [] + + glyph_texts = [extract_glyph_text(p) for p in prompt] + for glyph_text in glyph_texts: + if glyph_text is None: + glyph_text_embeds = torch.zeros( + (1, self.tokenizer_2_max_length, self.text_encoder_2.config.d_model), device=device + ) + glyph_text_embeds_mask = torch.zeros( + (1, self.tokenizer_2_max_length), device=device, dtype=torch.int64 + ) + else: + glyph_text_embeds, glyph_text_embeds_mask = self._get_byt5_prompt_embeds( + tokenizer=self.tokenizer_2, + text_encoder=self.text_encoder_2, + prompt=glyph_text, + device=device, + tokenizer_max_length=self.tokenizer_2_max_length, + ) + + prompt_embeds_2_list.append(glyph_text_embeds) + prompt_embeds_mask_2_list.append(glyph_text_embeds_mask) + + prompt_embeds_2 = torch.cat(prompt_embeds_2_list, dim=0) + prompt_embeds_mask_2 = torch.cat(prompt_embeds_mask_2_list, dim=0) + + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1) + prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len) + + _, seq_len_2, _ = prompt_embeds_2.shape + prompt_embeds_2 = prompt_embeds_2.repeat(1, num_images_per_prompt, 1) + prompt_embeds_2 = prompt_embeds_2.view(batch_size * num_images_per_prompt, seq_len_2, -1) + prompt_embeds_mask_2 = prompt_embeds_mask_2.repeat(1, num_images_per_prompt, 1) + prompt_embeds_mask_2 = prompt_embeds_mask_2.view(batch_size * num_images_per_prompt, seq_len_2) + + return prompt_embeds, prompt_embeds_mask, prompt_embeds_2, prompt_embeds_mask_2 + + def check_inputs( + self, + prompt, + height, + width, + negative_prompt=None, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_embeds_mask=None, + negative_prompt_embeds_mask=None, + prompt_embeds_2=None, + prompt_embeds_mask_2=None, + negative_prompt_embeds_2=None, + negative_prompt_embeds_mask_2=None, + callback_on_step_end_tensor_inputs=None, + ): + if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0: + logger.warning( + f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} but are {height} and {width}. Dimensions will be resized accordingly" + ) + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + if prompt_embeds is not None and prompt_embeds_mask is None: + raise ValueError( + "If `prompt_embeds` are provided, `prompt_embeds_mask` also have to be passed. Make sure to generate `prompt_embeds_mask` from the same text encoder that was used to generate `prompt_embeds`." + ) + if negative_prompt_embeds is not None and negative_prompt_embeds_mask is None: + raise ValueError( + "If `negative_prompt_embeds` are provided, `negative_prompt_embeds_mask` also have to be passed. Make sure to generate `negative_prompt_embeds_mask` from the same text encoder that was used to generate `negative_prompt_embeds`." + ) + + if prompt is None and prompt_embeds_2 is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds_2`. Cannot leave both `prompt` and `prompt_embeds_2` undefined." + ) + + if prompt_embeds_2 is not None and prompt_embeds_mask_2 is None: + raise ValueError( + "If `prompt_embeds_2` are provided, `prompt_embeds_mask_2` also have to be passed. Make sure to generate `prompt_embeds_mask_2` from the same text encoder that was used to generate `prompt_embeds_2`." + ) + if negative_prompt_embeds_2 is not None and negative_prompt_embeds_mask_2 is None: + raise ValueError( + "If `negative_prompt_embeds_2` are provided, `negative_prompt_embeds_mask_2` also have to be passed. Make sure to generate `negative_prompt_embeds_mask_2` from the same text encoder that was used to generate `negative_prompt_embeds_2`." + ) + + def prepare_latents( + self, + batch_size, + num_channels_latents, + height, + width, + dtype, + device, + generator, + latents=None, + ): + height = int(height) // self.vae_scale_factor + width = int(width) // self.vae_scale_factor + + shape = (batch_size, num_channels_latents, height, width) + + if latents is not None: + return latents.to(device=device, dtype=dtype) + + 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." + ) + + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + + return latents + + @property + def attention_kwargs(self): + return self._attention_kwargs + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: Union[str, List[str]] = None, + negative_prompt: Union[str, List[str]] = None, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 50, + distilled_guidance_scale: Optional[float] = 3.25, + sigmas: Optional[List[float]] = None, + num_images_per_prompt: int = 1, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds_mask: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt_embeds_mask: Optional[torch.Tensor] = None, + prompt_embeds_2: Optional[torch.Tensor] = None, + prompt_embeds_mask_2: Optional[torch.Tensor] = None, + negative_prompt_embeds_2: Optional[torch.Tensor] = None, + negative_prompt_embeds_mask_2: Optional[torch.Tensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + attention_kwargs: Optional[Dict[str, Any]] = None, + callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + ): + r""" + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined and negative_prompt_embeds is + not provided, will use an empty negative prompt. Ignored when not using guidance. ). + height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The height in pixels of the generated image. This is set to 1024 by default for the best results. + width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The width in pixels of the generated image. This is set to 1024 by default for the best results. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + sigmas (`List[float]`, *optional*): + Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in + their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed + will be used. + distilled_guidance_scale (`float`, *optional*, defaults to None): + A guidance scale value for guidance distilled models. Unlike the traditional classifier-free guidance + where the guidance scale is applied during inference through noise prediction rescaling, guidance + distilled models take the guidance scale directly as an input parameter during forward pass. Guidance + is enabled by setting `distilled_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. For + guidance distilled models, this parameter is required. For non-distilled models, this parameter will be + ignored. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will be generated by sampling using the supplied random `generator`. + prompt_embeds (`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. + prompt_embeds_mask (`torch.Tensor`, *optional*): + Pre-generated text embeddings mask. Can be used to easily tweak text inputs, *e.g.* prompt weighting. + If not provided, text embeddings mask will be generated from `prompt` input argument. + prompt_embeds_2 (`torch.Tensor`, *optional*): + Pre-generated text embeddings for ocr. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, text embeddings for ocr will be generated from `prompt` input argument. + prompt_embeds_mask_2 (`torch.Tensor`, *optional*): + Pre-generated text embeddings mask for ocr. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, text embeddings mask for ocr 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. + negative_prompt_embeds_mask (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings mask. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative text embeddings mask will be generated from `negative_prompt` + input argument. + negative_prompt_embeds_2 (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings for ocr. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative text embeddings for ocr will be generated from `negative_prompt` + input argument. + negative_prompt_embeds_mask_2 (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings mask for ocr. Can be used to easily tweak text inputs, *e.g.* + prompt weighting. If not provided, negative text embeddings mask for ocr will be generated from + `negative_prompt` input argument. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.qwenimage.QwenImagePipelineOutput`] instead of a plain tuple. + attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + callback_on_step_end (`Callable`, *optional*): + A function that calls at the end of each denoising steps during the inference. The function is called + with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, + callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by + `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + + Examples: + + Returns: + [`~pipelines.hunyuan_image.HunyuanImagePipelineOutput`] or `tuple`: + [`~pipelines.hunyuan_image.HunyuanImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When + returning a tuple, the first element is a list with the generated images. + """ + + height = height or self.default_sample_size * self.vae_scale_factor + width = width or self.default_sample_size * self.vae_scale_factor + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt, + height, + width, + negative_prompt=negative_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_embeds_mask=prompt_embeds_mask, + negative_prompt_embeds_mask=negative_prompt_embeds_mask, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + prompt_embeds_2=prompt_embeds_2, + prompt_embeds_mask_2=prompt_embeds_mask_2, + negative_prompt_embeds_2=negative_prompt_embeds_2, + negative_prompt_embeds_mask_2=negative_prompt_embeds_mask_2, + ) + + self._attention_kwargs = attention_kwargs + self._current_timestep = None + self._interrupt = False + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self._execution_device + + # 3. prepare prompt embeds + + prompt_embeds, prompt_embeds_mask, prompt_embeds_2, prompt_embeds_mask_2 = self.encode_prompt( + prompt=prompt, + prompt_embeds=prompt_embeds, + prompt_embeds_mask=prompt_embeds_mask, + device=device, + batch_size=batch_size, + num_images_per_prompt=num_images_per_prompt, + prompt_embeds_2=prompt_embeds_2, + prompt_embeds_mask_2=prompt_embeds_mask_2, + ) + + prompt_embeds = prompt_embeds.to(self.transformer.dtype) + prompt_embeds_2 = prompt_embeds_2.to(self.transformer.dtype) + + # select guider + if not torch.all(prompt_embeds_2 == 0) and self.ocr_guider is not None: + # prompt contains ocr and pipeline has a guider for ocr + guider = self.ocr_guider + elif self.guider is not None: + guider = self.guider + # distilled model does not use guidance method, use default guider with enabled=False + else: + guider = AdaptiveProjectedMixGuidance(enabled=False) + + if guider._enabled and guider.num_conditions > 1: + ( + negative_prompt_embeds, + negative_prompt_embeds_mask, + negative_prompt_embeds_2, + negative_prompt_embeds_mask_2, + ) = self.encode_prompt( + prompt=negative_prompt, + prompt_embeds=negative_prompt_embeds, + prompt_embeds_mask=negative_prompt_embeds_mask, + device=device, + batch_size=batch_size, + num_images_per_prompt=num_images_per_prompt, + prompt_embeds_2=negative_prompt_embeds_2, + prompt_embeds_mask_2=negative_prompt_embeds_mask_2, + ) + + negative_prompt_embeds = negative_prompt_embeds.to(self.transformer.dtype) + negative_prompt_embeds_2 = negative_prompt_embeds_2.to(self.transformer.dtype) + + # 4. Prepare latent variables + num_channels_latents = self.transformer.config.in_channels + latents = self.prepare_latents( + batch_size=batch_size * num_images_per_prompt, + num_channels_latents=num_channels_latents, + height=height, + width=width, + dtype=prompt_embeds.dtype, + device=device, + generator=generator, + latents=latents, + ) + + # 5. Prepare timesteps + sigmas = np.linspace(1.0, 0.0, num_inference_steps + 1)[:-1] if sigmas is None else sigmas + timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, sigmas=sigmas) + + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + + # handle guidance (for guidance-distilled model) + if self.transformer.config.guidance_embeds and distilled_guidance_scale is None: + raise ValueError("`distilled_guidance_scale` is required for guidance-distilled model.") + + if self.transformer.config.guidance_embeds: + guidance = ( + torch.tensor( + [distilled_guidance_scale] * latents.shape[0], dtype=self.transformer.dtype, device=device + ) + * 1000.0 + ) + + else: + guidance = None + + if self.attention_kwargs is None: + self._attention_kwargs = {} + + # 6. Denoising loop + self.scheduler.set_begin_index(0) + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + self._current_timestep = t + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latents.shape[0]).to(latents.dtype) + + if self.transformer.config.use_meanflow: + if i == len(timesteps) - 1: + timestep_r = torch.tensor([0.0], device=device) + else: + timestep_r = timesteps[i + 1] + timestep_r = timestep_r.expand(latents.shape[0]).to(latents.dtype) + else: + timestep_r = None + + # Step 1: Collect model inputs needed for the guidance method + # conditional inputs should always be first element in the tuple + guider_inputs = { + "encoder_hidden_states": (prompt_embeds, negative_prompt_embeds), + "encoder_attention_mask": (prompt_embeds_mask, negative_prompt_embeds_mask), + "encoder_hidden_states_2": (prompt_embeds_2, negative_prompt_embeds_2), + "encoder_attention_mask_2": (prompt_embeds_mask_2, negative_prompt_embeds_mask_2), + } + + # Step 2: Update guider's internal state for this denoising step + guider.set_state(step=i, num_inference_steps=num_inference_steps, timestep=t) + + # Step 3: Prepare batched model inputs based on the guidance method + # The guider splits model inputs into separate batches for conditional/unconditional predictions. + # For CFG with guider_inputs = {"encoder_hidden_states": (prompt_embeds, negative_prompt_embeds)}: + # you will get a guider_state with two batches: + # guider_state = [ + # {"encoder_hidden_states": prompt_embeds, "__guidance_identifier__": "pred_cond"}, # conditional batch + # {"encoder_hidden_states": negative_prompt_embeds, "__guidance_identifier__": "pred_uncond"}, # unconditional batch + # ] + # Other guidance methods may return 1 batch (no guidance) or 3+ batches (e.g., PAG, APG). + guider_state = guider.prepare_inputs(guider_inputs) + # Step 4: Run the denoiser for each batch + # Each batch in guider_state represents a different conditioning (conditional, unconditional, etc.). + # We run the model once per batch and store the noise prediction in guider_state_batch.noise_pred. + for guider_state_batch in guider_state: + guider.prepare_models(self.transformer) + + # Extract conditioning kwargs for this batch (e.g., encoder_hidden_states) + cond_kwargs = { + input_name: getattr(guider_state_batch, input_name) for input_name in guider_inputs.keys() + } + + # e.g. "pred_cond"/"pred_uncond" + context_name = getattr(guider_state_batch, guider._identifier_key) + with self.transformer.cache_context(context_name): + # Run denoiser and store noise prediction in this batch + guider_state_batch.noise_pred = self.transformer( + hidden_states=latents, + timestep=timestep, + timestep_r=timestep_r, + guidance=guidance, + attention_kwargs=self.attention_kwargs, + return_dict=False, + **cond_kwargs, + )[0] + + # Cleanup model (e.g., remove hooks) + guider.cleanup_models(self.transformer) + + # Step 5: Combine predictions using the guidance method + # The guider takes all noise predictions from guider_state and combines them according to the guidance algorithm. + # Continuing the CFG example, the guider receives: + # guider_state = [ + # {"encoder_hidden_states": prompt_embeds, "noise_pred": noise_pred_cond, "__guidance_identifier__": "pred_cond"}, # batch 0 + # {"encoder_hidden_states": negative_prompt_embeds, "noise_pred": noise_pred_uncond, "__guidance_identifier__": "pred_uncond"}, # batch 1 + # ] + # And extracts predictions using the __guidance_identifier__: + # pred_cond = guider_state[0]["noise_pred"] # extracts noise_pred_cond + # pred_uncond = guider_state[1]["noise_pred"] # extracts noise_pred_uncond + # Then applies CFG formula: + # noise_pred = pred_uncond + guidance_scale * (pred_cond - pred_uncond) + # Returns GuiderOutput(pred=noise_pred, pred_cond=pred_cond, pred_uncond=pred_uncond) + noise_pred = guider(guider_state)[0] + + # compute the previous noisy sample x_t -> x_t-1 + latents_dtype = latents.dtype + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + + if latents.dtype != latents_dtype: + if torch.backends.mps.is_available(): + # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272 + latents = latents.to(latents_dtype) + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + self._current_timestep = None + if output_type == "latent": + image = latents + else: + latents = latents.to(self.vae.dtype) / self.vae.config.scaling_factor + image = self.vae.decode(latents, return_dict=False)[0] + image = self.image_processor.postprocess(image, output_type=output_type) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return HunyuanImagePipelineOutput(images=image) diff --git a/src/diffusers/pipelines/hunyuan_image/pipeline_hunyuanimage_refiner.py b/src/diffusers/pipelines/hunyuan_image/pipeline_hunyuanimage_refiner.py new file mode 100644 index 0000000000..f38f53d9a5 --- /dev/null +++ b/src/diffusers/pipelines/hunyuan_image/pipeline_hunyuanimage_refiner.py @@ -0,0 +1,752 @@ +# Copyright 2025 Hunyuan-Image Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +from typing import Any, Callable, Dict, List, Optional, Union + +import numpy as np +import torch +from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer + +from ...guiders import AdaptiveProjectedMixGuidance +from ...image_processor import PipelineImageInput, VaeImageProcessor +from ...models import AutoencoderKLHunyuanImageRefiner, HunyuanImageTransformer2DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import is_torch_xla_available, logging, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ..pipeline_utils import DiffusionPipeline +from .pipeline_output import HunyuanImagePipelineOutput + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import HunyuanImageRefinerPipeline + + >>> pipe = HunyuanImageRefinerPipeline.from_pretrained( + ... "hunyuanvideo-community/HunyuanImage-2.1-Refiner-Diffusers", torch_dtype=torch.bfloat16 + ... ) + >>> pipe.to("cuda") + >>> prompt = "A cat holding a sign that says hello world" + >>> image = load_image("path/to/image.png") + >>> # Depending on the variant being used, the pipeline call will slightly vary. + >>> # Refer to the pipeline documentation for more details. + >>> image = pipe(prompt, image=image, num_inference_steps=4).images[0] + >>> image.save("hunyuanimage.png") + ``` +""" + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + timesteps: Optional[List[int]] = None, + sigmas: Optional[List[float]] = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`List[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`List[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents +def retrieve_latents( + encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" +): + if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": + return encoder_output.latent_dist.sample(generator) + elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": + return encoder_output.latent_dist.mode() + elif hasattr(encoder_output, "latents"): + return encoder_output.latents + else: + raise AttributeError("Could not access latents of provided encoder_output") + + +class HunyuanImageRefinerPipeline(DiffusionPipeline): + r""" + The HunyuanImage pipeline for text-to-image generation. + + Args: + transformer ([`HunyuanImageTransformer2DModel`]): + Conditional Transformer (MMDiT) architecture to denoise the encoded image latents. + scheduler ([`FlowMatchEulerDiscreteScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded image latents. + vae ([`AutoencoderKLHunyuanImageRefiner`]): + Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. + text_encoder ([`Qwen2.5-VL-7B-Instruct`]): + [Qwen2.5-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct), specifically the + [Qwen2.5-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct) variant. + tokenizer (`Qwen2Tokenizer`): Tokenizer of class [Qwen2Tokenizer]. + """ + + model_cpu_offload_seq = "text_encoder->transformer->vae" + _callback_tensor_inputs = ["latents", "prompt_embeds"] + _optional_components = ["guider"] + + def __init__( + self, + scheduler: FlowMatchEulerDiscreteScheduler, + vae: AutoencoderKLHunyuanImageRefiner, + text_encoder: Qwen2_5_VLForConditionalGeneration, + tokenizer: Qwen2Tokenizer, + transformer: HunyuanImageTransformer2DModel, + guider: Optional[AdaptiveProjectedMixGuidance] = None, + ): + super().__init__() + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + scheduler=scheduler, + guider=guider, + ) + + self.vae_scale_factor = self.vae.config.spatial_compression_ratio if getattr(self, "vae", None) else 16 + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + self.tokenizer_max_length = 256 + self.prompt_template_encode = "<|start_header_id|>system<|end_header_id|>\n\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>" + self.prompt_template_encode_start_idx = 36 + self.default_sample_size = 64 + self.latent_channels = self.transformer.config.in_channels // 2 if getattr(self, "transformer", None) else 64 + + # Copied from diffusers.pipelines.hunyuan_image.pipeline_hunyuanimage.HunyuanImagePipeline._get_qwen_prompt_embeds + def _get_qwen_prompt_embeds( + self, + tokenizer: Qwen2Tokenizer, + text_encoder: Qwen2_5_VLForConditionalGeneration, + prompt: Union[str, List[str]] = None, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + tokenizer_max_length: int = 1000, + template: str = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>", + drop_idx: int = 34, + hidden_state_skip_layer: int = 2, + ): + device = device or self._execution_device + dtype = dtype or text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + + txt = [template.format(e) for e in prompt] + txt_tokens = tokenizer( + txt, max_length=tokenizer_max_length + drop_idx, padding="max_length", truncation=True, return_tensors="pt" + ).to(device) + + encoder_hidden_states = text_encoder( + input_ids=txt_tokens.input_ids, + attention_mask=txt_tokens.attention_mask, + output_hidden_states=True, + ) + prompt_embeds = encoder_hidden_states.hidden_states[-(hidden_state_skip_layer + 1)] + + prompt_embeds = prompt_embeds[:, drop_idx:] + encoder_attention_mask = txt_tokens.attention_mask[:, drop_idx:] + + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + encoder_attention_mask = encoder_attention_mask.to(device=device) + + return prompt_embeds, encoder_attention_mask + + def encode_prompt( + self, + prompt: Optional[Union[str, List[str]]] = None, + device: Optional[torch.device] = None, + batch_size: int = 1, + num_images_per_prompt: int = 1, + prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds_mask: Optional[torch.Tensor] = None, + ): + r""" + + Args: + prompt (`str` or `List[str]`, *optional*): + prompt to be encoded + device: (`torch.device`): + torch device + batch_size (`int`): + batch size of prompts, defaults to 1 + num_images_per_prompt (`int`): + number of images that should be generated per prompt + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. If not provided, text embeddings will be generated from `prompt` input + argument. + prompt_embeds_mask (`torch.Tensor`, *optional*): + Pre-generated text mask. If not provided, text mask will be generated from `prompt` input argument. + prompt_embeds_2 (`torch.Tensor`, *optional*): + Pre-generated glyph text embeddings from ByT5. If not provided, will be generated from `prompt` input + argument using self.tokenizer_2 and self.text_encoder_2. + prompt_embeds_mask_2 (`torch.Tensor`, *optional*): + Pre-generated glyph text mask from ByT5. If not provided, will be generated from `prompt` input + argument using self.tokenizer_2 and self.text_encoder_2. + """ + device = device or self._execution_device + + if prompt is None: + prompt = [""] * batch_size + + prompt = [prompt] if isinstance(prompt, str) else prompt + + if prompt_embeds is None: + prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds( + tokenizer=self.tokenizer, + text_encoder=self.text_encoder, + prompt=prompt, + device=device, + tokenizer_max_length=self.tokenizer_max_length, + template=self.prompt_template_encode, + drop_idx=self.prompt_template_encode_start_idx, + ) + + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1) + prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len) + + return prompt_embeds, prompt_embeds_mask + + def check_inputs( + self, + prompt, + height, + width, + negative_prompt=None, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_embeds_mask=None, + negative_prompt_embeds_mask=None, + callback_on_step_end_tensor_inputs=None, + ): + if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0: + logger.warning( + f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} but are {height} and {width}. Dimensions will be resized accordingly" + ) + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + if prompt_embeds is not None and prompt_embeds_mask is None: + raise ValueError( + "If `prompt_embeds` are provided, `prompt_embeds_mask` also have to be passed. Make sure to generate `prompt_embeds_mask` from the same text encoder that was used to generate `prompt_embeds`." + ) + if negative_prompt_embeds is not None and negative_prompt_embeds_mask is None: + raise ValueError( + "If `negative_prompt_embeds` are provided, `negative_prompt_embeds_mask` also have to be passed. Make sure to generate `negative_prompt_embeds_mask` from the same text encoder that was used to generate `negative_prompt_embeds`." + ) + + def prepare_latents( + self, + image_latents, + batch_size, + num_channels_latents, + height, + width, + dtype, + device, + generator, + latents=None, + strength=0.25, + ): + height = int(height) // self.vae_scale_factor + width = int(width) // self.vae_scale_factor + + shape = (batch_size, num_channels_latents, 1, height, width) + + if latents is None: + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + else: + latents = latents.to(device=device, dtype=dtype) + + if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0: + # expand init_latents for batch_size + additional_image_per_prompt = batch_size // image_latents.shape[0] + image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0) + elif batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0: + raise ValueError( + f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts." + ) + + 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." + ) + + noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + cond_latents = strength * noise + (1 - strength) * image_latents + + return latents, cond_latents + + @staticmethod + def _reorder_image_tokens(image_latents): + image_latents = torch.cat((image_latents[:, :, :1], image_latents), dim=2) + batch_size, num_latent_channels, num_latent_frames, latent_height, latent_width = image_latents.shape + image_latents = image_latents.permute(0, 2, 1, 3, 4) + image_latents = image_latents.reshape( + batch_size, num_latent_frames // 2, num_latent_channels * 2, latent_height, latent_width + ) + image_latents = image_latents.permute(0, 2, 1, 3, 4).contiguous() + + return image_latents + + @staticmethod + def _restore_image_tokens_order(latents): + """Restore image tokens order by splitting channels and removing first frame slice.""" + batch_size, num_latent_channels, num_latent_frames, latent_height, latent_width = latents.shape + + latents = latents.permute(0, 2, 1, 3, 4) # B, F, C, H, W + latents = latents.reshape( + batch_size, num_latent_frames * 2, num_latent_channels // 2, latent_height, latent_width + ) # B, F*2, C//2, H, W + + latents = latents.permute(0, 2, 1, 3, 4) # B, C//2, F*2, H, W + # Remove first frame slice + latents = latents[:, :, 1:] + + return latents + + def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): + if isinstance(generator, list): + image_latents = [ + retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i], sample_mode="sample") + 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, sample_mode="sample") + image_latents = self._reorder_image_tokens(image_latents) + + image_latents = image_latents * self.vae.config.scaling_factor + + return image_latents + + @property + def attention_kwargs(self): + return self._attention_kwargs + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: Union[str, List[str]] = None, + negative_prompt: Union[str, List[str]] = None, + distilled_guidance_scale: Optional[float] = 3.25, + image: Optional[PipelineImageInput] = None, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 4, + sigmas: Optional[List[float]] = None, + num_images_per_prompt: int = 1, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds_mask: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt_embeds_mask: Optional[torch.Tensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + attention_kwargs: Optional[Dict[str, Any]] = None, + callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + ): + r""" + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, will use an empty negative + prompt. Ignored when not using guidance. + distilled_guidance_scale (`float`, *optional*, defaults to None): + A guidance scale value for guidance distilled models. Unlike the traditional classifier-free guidance + where the guidance scale is applied during inference through noise prediction rescaling, guidance + distilled models take the guidance scale directly as an input parameter during forward pass. Guidance + is enabled by setting `distilled_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. For + guidance distilled models, this parameter is required. For non-distilled models, this parameter will be + ignored. + num_images_per_prompt (`int`, *optional*, defaults to 1): + height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The height in pixels of the generated image. This is set to 1024 by default for the best results. + width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The width in pixels of the generated image. This is set to 1024 by default for the best results. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + sigmas (`List[float]`, *optional*): + Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in + their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed + will be used. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will be generated by sampling using the supplied random `generator`. + prompt_embeds (`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. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.qwenimage.QwenImagePipelineOutput`] instead of a plain tuple. + attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + callback_on_step_end (`Callable`, *optional*): + A function that calls at the end of each denoising steps during the inference. The function is called + with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, + callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by + `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + + Examples: + + Returns: + [`~pipelines.hunyuan_image.HunyuanImagePipelineOutput`] or `tuple`: + [`~pipelines.hunyuan_image.HunyuanImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When + returning a tuple, the first element is a list with the generated images. + """ + + height = height or self.default_sample_size * self.vae_scale_factor + width = width or self.default_sample_size * self.vae_scale_factor + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt, + height, + width, + negative_prompt=negative_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_embeds_mask=prompt_embeds_mask, + negative_prompt_embeds_mask=negative_prompt_embeds_mask, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + ) + + self._attention_kwargs = attention_kwargs + self._current_timestep = None + self._interrupt = False + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self._execution_device + + # 3. process image + if image is not None and isinstance(image, torch.Tensor) and image.shape[1] == self.latent_channels: + image_latents = image + else: + image = self.image_processor.preprocess(image, height, width) + image = image.unsqueeze(2).to(device, dtype=self.vae.dtype) + image_latents = self._encode_vae_image(image=image, generator=generator) + + # 3.prepare prompt embeds + + if self.guider is not None: + guider = self.guider + else: + # distilled model does not use guidance method, use default guider with enabled=False + guider = AdaptiveProjectedMixGuidance(enabled=False) + + requires_unconditional_embeds = guider._enabled and guider.num_conditions > 1 + prompt_embeds, prompt_embeds_mask = self.encode_prompt( + prompt=prompt, + prompt_embeds=prompt_embeds, + prompt_embeds_mask=prompt_embeds_mask, + device=device, + batch_size=batch_size, + num_images_per_prompt=num_images_per_prompt, + ) + + prompt_embeds = prompt_embeds.to(self.transformer.dtype) + + if requires_unconditional_embeds: + ( + negative_prompt_embeds, + negative_prompt_embeds_mask, + ) = self.encode_prompt( + prompt=negative_prompt, + prompt_embeds=negative_prompt_embeds, + prompt_embeds_mask=negative_prompt_embeds_mask, + device=device, + batch_size=batch_size, + num_images_per_prompt=num_images_per_prompt, + ) + + negative_prompt_embeds = negative_prompt_embeds.to(self.transformer.dtype) + + # 4. Prepare latent variables + latents, cond_latents = self.prepare_latents( + image_latents=image_latents, + batch_size=batch_size * num_images_per_prompt, + num_channels_latents=self.latent_channels, + height=height, + width=width, + dtype=prompt_embeds.dtype, + device=device, + generator=generator, + latents=latents, + ) + + # 5. Prepare timesteps + sigmas = np.linspace(1.0, 0.0, num_inference_steps + 1)[:-1] if sigmas is None else sigmas + timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, sigmas=sigmas) + + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + + # handle guidance (this pipeline only supports guidance-distilled models) + if distilled_guidance_scale is None: + raise ValueError("`distilled_guidance_scale` is required for guidance-distilled model.") + guidance = ( + torch.tensor([distilled_guidance_scale] * latents.shape[0], dtype=self.transformer.dtype, device=device) + * 1000.0 + ) + + if self.attention_kwargs is None: + self._attention_kwargs = {} + + # 6. Denoising loop + self.scheduler.set_begin_index(0) + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + self._current_timestep = t + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + latent_model_input = torch.cat([latents, cond_latents], dim=1).to(self.transformer.dtype) + timestep = t.expand(latents.shape[0]).to(latents.dtype) + + # Step 1: Collect model inputs needed for the guidance method + # conditional inputs should always be first element in the tuple + guider_inputs = { + "encoder_hidden_states": (prompt_embeds, negative_prompt_embeds), + "encoder_attention_mask": (prompt_embeds_mask, negative_prompt_embeds_mask), + } + + # Step 2: Update guider's internal state for this denoising step + guider.set_state(step=i, num_inference_steps=num_inference_steps, timestep=t) + + # Step 3: Prepare batched model inputs based on the guidance method + # The guider splits model inputs into separate batches for conditional/unconditional predictions. + # For CFG with guider_inputs = {"encoder_hidden_states": (prompt_embeds, negative_prompt_embeds)}: + # you will get a guider_state with two batches: + # guider_state = [ + # {"encoder_hidden_states": prompt_embeds, "__guidance_identifier__": "pred_cond"}, # conditional batch + # {"encoder_hidden_states": negative_prompt_embeds, "__guidance_identifier__": "pred_uncond"}, # unconditional batch + # ] + # Other guidance methods may return 1 batch (no guidance) or 3+ batches (e.g., PAG, APG). + guider_state = guider.prepare_inputs(guider_inputs) + + # Step 4: Run the denoiser for each batch + # Each batch in guider_state represents a different conditioning (conditional, unconditional, etc.). + # We run the model once per batch and store the noise prediction in guider_state_batch.noise_pred. + for guider_state_batch in guider_state: + guider.prepare_models(self.transformer) + + # Extract conditioning kwargs for this batch (e.g., encoder_hidden_states) + cond_kwargs = { + input_name: getattr(guider_state_batch, input_name) for input_name in guider_inputs.keys() + } + + # e.g. "pred_cond"/"pred_uncond" + context_name = getattr(guider_state_batch, guider._identifier_key) + with self.transformer.cache_context(context_name): + # Run denoiser and store noise prediction in this batch + guider_state_batch.noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + guidance=guidance, + attention_kwargs=self.attention_kwargs, + return_dict=False, + **cond_kwargs, + )[0] + + # Cleanup model (e.g., remove hooks) + guider.cleanup_models(self.transformer) + + # Step 5: Combine predictions using the guidance method + # The guider takes all noise predictions from guider_state and combines them according to the guidance algorithm. + # Continuing the CFG example, the guider receives: + # guider_state = [ + # {"encoder_hidden_states": prompt_embeds, "noise_pred": noise_pred_cond, "__guidance_identifier__": "pred_cond"}, # batch 0 + # {"encoder_hidden_states": negative_prompt_embeds, "noise_pred": noise_pred_uncond, "__guidance_identifier__": "pred_uncond"}, # batch 1 + # ] + # And extracts predictions using the __guidance_identifier__: + # pred_cond = guider_state[0]["noise_pred"] # extracts noise_pred_cond + # pred_uncond = guider_state[1]["noise_pred"] # extracts noise_pred_uncond + # Then applies CFG formula: + # noise_pred = pred_uncond + guidance_scale * (pred_cond - pred_uncond) + # Returns GuiderOutput(pred=noise_pred, pred_cond=pred_cond, pred_uncond=pred_uncond) + noise_pred = guider(guider_state)[0] + + # compute the previous noisy sample x_t -> x_t-1 + latents_dtype = latents.dtype + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + + if latents.dtype != latents_dtype: + if torch.backends.mps.is_available(): + # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272 + latents = latents.to(latents_dtype) + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + self._current_timestep = None + if output_type == "latent": + image = latents + else: + latents = latents.to(self.vae.dtype) / self.vae.config.scaling_factor + latents = self._restore_image_tokens_order(latents) + + image = self.vae.decode(latents, return_dict=False)[0] + image = self.image_processor.postprocess(image.squeeze(2), output_type=output_type) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return HunyuanImagePipelineOutput(images=image) diff --git a/src/diffusers/pipelines/hunyuan_image/pipeline_output.py b/src/diffusers/pipelines/hunyuan_image/pipeline_output.py new file mode 100644 index 0000000000..1e76892a0e --- /dev/null +++ b/src/diffusers/pipelines/hunyuan_image/pipeline_output.py @@ -0,0 +1,21 @@ +from dataclasses import dataclass +from typing import List, Union + +import numpy as np +import PIL.Image + +from ...utils import BaseOutput + + +@dataclass +class HunyuanImagePipelineOutput(BaseOutput): + """ + Output class for HunyuanImage pipelines. + + Args: + images (`List[PIL.Image.Image]` or `np.ndarray`) + List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, + num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. + """ + + images: Union[List[PIL.Image.Image], np.ndarray] diff --git a/src/diffusers/pipelines/pipeline_loading_utils.py b/src/diffusers/pipelines/pipeline_loading_utils.py index 00ef996185..8868e942ce 100644 --- a/src/diffusers/pipelines/pipeline_loading_utils.py +++ b/src/diffusers/pipelines/pipeline_loading_utils.py @@ -76,6 +76,7 @@ LOADABLE_CLASSES = { "SchedulerMixin": ["save_pretrained", "from_pretrained"], "DiffusionPipeline": ["save_pretrained", "from_pretrained"], "OnnxRuntimeModel": ["save_pretrained", "from_pretrained"], + "BaseGuidance": ["save_pretrained", "from_pretrained"], }, "transformers": { "PreTrainedTokenizer": ["save_pretrained", "from_pretrained"], diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index ecf2d7957a..6ef6b8b0e9 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -17,6 +17,21 @@ class AdaptiveProjectedGuidance(metaclass=DummyObject): requires_backends(cls, ["torch"]) +class AdaptiveProjectedMixGuidance(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class AutoGuidance(metaclass=DummyObject): _backends = ["torch"] @@ -32,6 +47,21 @@ class AutoGuidance(metaclass=DummyObject): requires_backends(cls, ["torch"]) +class BaseGuidance(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class ClassifierFreeGuidance(metaclass=DummyObject): _backends = ["torch"] @@ -378,6 +408,36 @@ class AutoencoderKLCosmos(metaclass=DummyObject): requires_backends(cls, ["torch"]) +class AutoencoderKLHunyuanImage(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + +class AutoencoderKLHunyuanImageRefiner(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class AutoencoderKLHunyuanVideo(metaclass=DummyObject): _backends = ["torch"] @@ -858,6 +918,21 @@ class HunyuanDiT2DMultiControlNetModel(metaclass=DummyObject): requires_backends(cls, ["torch"]) +class HunyuanImageTransformer2DModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class HunyuanVideoFramepackTransformer3DModel(metaclass=DummyObject): _backends = ["torch"] diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 3a106c1b83..7c1dcba9c7 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -1037,6 +1037,36 @@ class HunyuanDiTPipeline(metaclass=DummyObject): requires_backends(cls, ["torch", "transformers"]) +class HunyuanImagePipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class HunyuanImageRefinerPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class HunyuanSkyreelsImageToVideoPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/tests/pipelines/hunyuan_image_21/__init__.py b/tests/pipelines/hunyuan_image_21/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/pipelines/hunyuan_image_21/test_hunyuanimage.py b/tests/pipelines/hunyuan_image_21/test_hunyuanimage.py new file mode 100644 index 0000000000..e4b2c686b8 --- /dev/null +++ b/tests/pipelines/hunyuan_image_21/test_hunyuanimage.py @@ -0,0 +1,290 @@ +# Copyright 2025 The HuggingFace Team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import numpy as np +import torch +from transformers import ( + ByT5Tokenizer, + Qwen2_5_VLConfig, + Qwen2_5_VLForConditionalGeneration, + Qwen2Tokenizer, + T5Config, + T5EncoderModel, +) + +from diffusers import ( + AdaptiveProjectedMixGuidance, + AutoencoderKLHunyuanImage, + FlowMatchEulerDiscreteScheduler, + HunyuanImagePipeline, + HunyuanImageTransformer2DModel, +) + +from ...testing_utils import enable_full_determinism +from ..test_pipelines_common import ( + FirstBlockCacheTesterMixin, + PipelineTesterMixin, + to_np, +) + + +enable_full_determinism() + + +class HunyuanImagePipelineFastTests( + PipelineTesterMixin, + FirstBlockCacheTesterMixin, + unittest.TestCase, +): + pipeline_class = HunyuanImagePipeline + params = frozenset(["prompt", "height", "width"]) + batch_params = frozenset(["prompt", "negative_prompt"]) + required_optional_params = frozenset( + [ + "num_inference_steps", + "generator", + "latents", + "return_dict", + "callback_on_step_end", + "callback_on_step_end_tensor_inputs", + ] + ) + + test_xformers_attention = False + test_layerwise_casting = True + test_group_offloading = True + test_attention_slicing = False + supports_dduf = False + + def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1, guidance_embeds: bool = False): + torch.manual_seed(0) + transformer = HunyuanImageTransformer2DModel( + in_channels=4, + out_channels=4, + num_attention_heads=4, + attention_head_dim=8, + num_layers=num_layers, + num_single_layers=num_single_layers, + num_refiner_layers=1, + patch_size=(1, 1), + guidance_embeds=guidance_embeds, + text_embed_dim=32, + text_embed_2_dim=32, + rope_axes_dim=(4, 4), + ) + + torch.manual_seed(0) + vae = AutoencoderKLHunyuanImage( + in_channels=3, + out_channels=3, + latent_channels=4, + block_out_channels=(32, 64, 64, 64), + layers_per_block=1, + scaling_factor=0.476986, + spatial_compression_ratio=8, + sample_size=128, + ) + + torch.manual_seed(0) + scheduler = FlowMatchEulerDiscreteScheduler(shift=7.0) + + if not guidance_embeds: + torch.manual_seed(0) + guider = AdaptiveProjectedMixGuidance(adaptive_projected_guidance_start_step=2) + ocr_guider = AdaptiveProjectedMixGuidance(adaptive_projected_guidance_start_step=3) + else: + guider = None + ocr_guider = None + torch.manual_seed(0) + config = Qwen2_5_VLConfig( + text_config={ + "hidden_size": 32, + "intermediate_size": 32, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 2, + "rope_scaling": { + "mrope_section": [2, 2, 4], + "rope_type": "default", + "type": "default", + }, + "rope_theta": 1000000.0, + }, + vision_config={ + "depth": 2, + "hidden_size": 32, + "intermediate_size": 32, + "num_heads": 2, + "out_hidden_size": 32, + }, + hidden_size=32, + vocab_size=152064, + vision_end_token_id=151653, + vision_start_token_id=151652, + vision_token_id=151654, + ) + text_encoder = Qwen2_5_VLForConditionalGeneration(config) + tokenizer = Qwen2Tokenizer.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration") + + torch.manual_seed(0) + t5_config = T5Config( + d_model=32, + d_kv=4, + d_ff=16, + num_layers=2, + num_heads=2, + relative_attention_num_buckets=8, + relative_attention_max_distance=32, + vocab_size=256, + feed_forward_proj="gated-gelu", + dense_act_fn="gelu_new", + is_encoder_decoder=False, + use_cache=False, + tie_word_embeddings=False, + ) + text_encoder_2 = T5EncoderModel(t5_config) + tokenizer_2 = ByT5Tokenizer() + + components = { + "transformer": transformer, + "vae": vae, + "scheduler": scheduler, + "text_encoder": text_encoder, + "text_encoder_2": text_encoder_2, + "tokenizer": tokenizer, + "tokenizer_2": tokenizer_2, + "guider": guider, + "ocr_guider": ocr_guider, + } + return components + + def get_dummy_inputs(self, device, seed=0): + if str(device).startswith("mps"): + generator = torch.manual_seed(seed) + else: + generator = torch.Generator(device=device).manual_seed(seed) + + inputs = { + "prompt": "A painting of a squirrel eating a burger", + "generator": generator, + "num_inference_steps": 5, + "height": 16, + "width": 16, + "output_type": "pt", + } + return inputs + + def test_inference(self): + device = "cpu" + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + image = pipe(**inputs).images + generated_image = image[0] + self.assertEqual(generated_image.shape, (3, 16, 16)) + + expected_slice_np = np.array( + [0.6252659, 0.51482046, 0.60799813, 0.59267783, 0.488082, 0.5857634, 0.523781, 0.58028054, 0.5674121] + ) + output_slice = generated_image[0, -3:, -3:].flatten().cpu().numpy() + + self.assertTrue( + np.abs(output_slice - expected_slice_np).max() < 1e-3, + f"output_slice: {output_slice}, expected_slice_np: {expected_slice_np}", + ) + + def test_inference_guider(self): + device = "cpu" + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + pipe.guider = pipe.guider.new(guidance_scale=1000) + pipe.ocr_guider = pipe.ocr_guider.new(guidance_scale=1000) + + inputs = self.get_dummy_inputs(device) + image = pipe(**inputs).images + generated_image = image[0] + self.assertEqual(generated_image.shape, (3, 16, 16)) + + expected_slice_np = np.array( + [0.61494756, 0.49616697, 0.60327923, 0.6115793, 0.49047345, 0.56977504, 0.53066164, 0.58880305, 0.5570612] + ) + output_slice = generated_image[0, -3:, -3:].flatten().cpu().numpy() + + self.assertTrue( + np.abs(output_slice - expected_slice_np).max() < 1e-3, + f"output_slice: {output_slice}, expected_slice_np: {expected_slice_np}", + ) + + def test_inference_with_distilled_guidance(self): + device = "cpu" + + components = self.get_dummy_components(guidance_embeds=True) + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + inputs["distilled_guidance_scale"] = 3.5 + image = pipe(**inputs).images + generated_image = image[0] + self.assertEqual(generated_image.shape, (3, 16, 16)) + + expected_slice_np = np.array( + [0.63667065, 0.5187377, 0.66757566, 0.6320319, 0.4913387, 0.54813194, 0.5335031, 0.5736143, 0.5461346] + ) + output_slice = generated_image[0, -3:, -3:].flatten().cpu().numpy() + + self.assertTrue( + np.abs(output_slice - expected_slice_np).max() < 1e-3, + f"output_slice: {output_slice}, expected_slice_np: {expected_slice_np}", + ) + + def test_vae_tiling(self, expected_diff_max: float = 0.2): + generator_device = "cpu" + components = self.get_dummy_components() + + pipe = self.pipeline_class(**components) + pipe.to("cpu") + pipe.set_progress_bar_config(disable=None) + + # Without tiling + inputs = self.get_dummy_inputs(generator_device) + inputs["height"] = inputs["width"] = 128 + output_without_tiling = pipe(**inputs)[0] + + # With tiling + pipe.vae.enable_tiling(tile_sample_min_size=96) + inputs = self.get_dummy_inputs(generator_device) + inputs["height"] = inputs["width"] = 128 + output_with_tiling = pipe(**inputs)[0] + + self.assertLess( + (to_np(output_without_tiling) - to_np(output_with_tiling)).max(), + expected_diff_max, + "VAE tiling should not affect the inference results", + ) + + @unittest.skip("TODO: Test not supported for now because needs to be adjusted to work with guiders.") + def test_encode_prompt_works_in_isolation(self): + pass From 7536f647e4144c7acaf9e140893ff7edb85bf9a3 Mon Sep 17 00:00:00 2001 From: kaixuanliu Date: Fri, 24 Oct 2025 18:18:15 +0800 Subject: [PATCH 093/123] Loose the criteria tolerance appropriately for Intel XPU devices (#12460) * Loose the criteria tolerance appropriately for Intel XPU devices Signed-off-by: Liu, Kaixuan * change back the atol value Signed-off-by: Liu, Kaixuan * use expectations Signed-off-by: Liu, Kaixuan * Update tests/pipelines/kandinsky2_2/test_kandinsky_controlnet.py --------- Signed-off-by: Liu, Kaixuan Co-authored-by: Ilyas Moutawwakil <57442720+IlyasMoutawwakil@users.noreply.github.com> --- .../kandinsky2_2/test_kandinsky_controlnet.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/pipelines/kandinsky2_2/test_kandinsky_controlnet.py b/tests/pipelines/kandinsky2_2/test_kandinsky_controlnet.py index 4054e38c56..8f8e58a8c4 100644 --- a/tests/pipelines/kandinsky2_2/test_kandinsky_controlnet.py +++ b/tests/pipelines/kandinsky2_2/test_kandinsky_controlnet.py @@ -29,6 +29,7 @@ from diffusers import ( ) from ...testing_utils import ( + Expectations, backend_empty_cache, enable_full_determinism, floats_tensor, @@ -290,4 +291,11 @@ class KandinskyV22ControlnetPipelineIntegrationTests(unittest.TestCase): assert image.shape == (512, 512, 3) max_diff = numpy_cosine_similarity_distance(expected_image.flatten(), image.flatten()) - assert max_diff < 2e-4 + expected_max_diffs = Expectations( + { + ("xpu", 3): 2e-3, + ("cuda", 7): 2e-4, + } + ) + expected_max_diff = expected_max_diffs.get_expectation() + assert max_diff < expected_max_diff From d34b18c783a772a0f9b24b7da7b834d785b5e58f Mon Sep 17 00:00:00 2001 From: Dhruv Nair Date: Fri, 24 Oct 2025 22:06:31 +0530 Subject: [PATCH 094/123] Deprecate Stable Cascade (#12537) * update * update --------- Co-authored-by: Sayak Paul --- .../pipelines/stable_cascade/pipeline_stable_cascade.py | 6 ++++-- .../stable_cascade/pipeline_stable_cascade_combined.py | 6 ++++-- .../stable_cascade/pipeline_stable_cascade_prior.py | 6 ++++-- 3 files changed, 12 insertions(+), 6 deletions(-) diff --git a/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py b/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py index aa39983c4e..a6a60ad94b 100644 --- a/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py +++ b/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py @@ -21,7 +21,7 @@ from ...models import StableCascadeUNet from ...schedulers import DDPMWuerstchenScheduler from ...utils import is_torch_version, is_torch_xla_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput +from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline, ImagePipelineOutput from ..wuerstchen.modeling_paella_vq_model import PaellaVQModel @@ -55,7 +55,7 @@ EXAMPLE_DOC_STRING = """ """ -class StableCascadeDecoderPipeline(DiffusionPipeline): +class StableCascadeDecoderPipeline(DeprecatedPipelineMixin, DiffusionPipeline): """ Pipeline for generating images from the Stable Cascade model. @@ -79,6 +79,8 @@ class StableCascadeDecoderPipeline(DiffusionPipeline): width=int(24*10.67)=256 in order to match the training conditions. """ + _last_supported_version = "0.35.2" + unet_name = "decoder" text_encoder_name = "text_encoder" model_cpu_offload_seq = "text_encoder->decoder->vqgan" diff --git a/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_combined.py b/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_combined.py index b3dc23f2e5..838b93faaa 100644 --- a/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_combined.py +++ b/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_combined.py @@ -20,7 +20,7 @@ from transformers import CLIPImageProcessor, CLIPTextModelWithProjection, CLIPTo from ...models import StableCascadeUNet from ...schedulers import DDPMWuerstchenScheduler from ...utils import is_torch_version, replace_example_docstring -from ..pipeline_utils import DiffusionPipeline +from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline from ..wuerstchen.modeling_paella_vq_model import PaellaVQModel from .pipeline_stable_cascade import StableCascadeDecoderPipeline from .pipeline_stable_cascade_prior import StableCascadePriorPipeline @@ -42,7 +42,7 @@ TEXT2IMAGE_EXAMPLE_DOC_STRING = """ """ -class StableCascadeCombinedPipeline(DiffusionPipeline): +class StableCascadeCombinedPipeline(DeprecatedPipelineMixin, DiffusionPipeline): """ Combined Pipeline for text-to-image generation using Stable Cascade. @@ -74,6 +74,8 @@ class StableCascadeCombinedPipeline(DiffusionPipeline): Frozen CLIP image-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). """ + _last_supported_version = "0.35.2" + _load_connected_pipes = True _optional_components = ["prior_feature_extractor", "prior_image_encoder"] diff --git a/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_prior.py b/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_prior.py index 9e63b3489c..29ad8b5429 100644 --- a/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_prior.py +++ b/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_prior.py @@ -25,7 +25,7 @@ from ...models import StableCascadeUNet from ...schedulers import DDPMWuerstchenScheduler from ...utils import BaseOutput, is_torch_xla_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline +from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline if is_torch_xla_available(): @@ -77,7 +77,7 @@ class StableCascadePriorPipelineOutput(BaseOutput): negative_prompt_embeds_pooled: Union[torch.Tensor, np.ndarray] -class StableCascadePriorPipeline(DiffusionPipeline): +class StableCascadePriorPipeline(DeprecatedPipelineMixin, DiffusionPipeline): """ Pipeline for generating image prior for Stable Cascade. @@ -103,6 +103,8 @@ class StableCascadePriorPipeline(DiffusionPipeline): Default resolution for multiple images generated. """ + _last_supported_version = "0.35.2" + unet_name = "prior" text_encoder_name = "text_encoder" model_cpu_offload_seq = "image_encoder->text_encoder->prior" From 500b9cf184d9a6aa950ccbb9a4d967b95877196d Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Sun, 26 Oct 2025 23:11:23 +0530 Subject: [PATCH 095/123] [chore] Move guiders experimental warning (#12543) * move guiders experimental warning to init. * up --- src/diffusers/guiders/__init__.py | 6 ------ src/diffusers/guiders/guider_utils.py | 4 ++++ 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/src/diffusers/guiders/__init__.py b/src/diffusers/guiders/__init__.py index 6c95b8e0c4..4e53c373c4 100644 --- a/src/diffusers/guiders/__init__.py +++ b/src/diffusers/guiders/__init__.py @@ -17,12 +17,6 @@ from typing import Union from ..utils import is_torch_available, logging -logger = logging.get_logger(__name__) -logger.warning( - "Guiders are currently an experimental feature under active development. The API is subject to breaking changes in future releases." -) - - if is_torch_available(): from .adaptive_projected_guidance import AdaptiveProjectedGuidance from .adaptive_projected_guidance_mix import AdaptiveProjectedMixGuidance diff --git a/src/diffusers/guiders/guider_utils.py b/src/diffusers/guiders/guider_utils.py index b7afd8648b..71e4becfcd 100644 --- a/src/diffusers/guiders/guider_utils.py +++ b/src/diffusers/guiders/guider_utils.py @@ -41,6 +41,10 @@ class BaseGuidance(ConfigMixin, PushToHubMixin): _identifier_key = "__guidance_identifier__" def __init__(self, start: float = 0.0, stop: float = 1.0, enabled: bool = True): + logger.warning( + "Guiders are currently an experimental feature under active development. The API is subject to breaking changes in future releases." + ) + self._start = start self._stop = stop self._step: int = None From dc6bd1511a4948ebca35b22609002bba58e71c83 Mon Sep 17 00:00:00 2001 From: josephrocca <1167575+josephrocca@users.noreply.github.com> Date: Mon, 27 Oct 2025 18:55:20 +0800 Subject: [PATCH 096/123] Fix Chroma attention padding order and update docs to use `lodestones/Chroma1-HD` (#12508) * [Fix] Move attention mask padding after T5 embedding * [Fix] Move attention mask padding after T5 embedding * Clean up whitespace in pipeline_chroma.py Removed unnecessary blank lines for cleaner code. * Fix * Fix * Update model to final Chroma1-HD checkpoint * Update to Chroma1-HD * Update model to Chroma1-HD * Update model to Chroma1-HD * Update Chroma model links to Chroma1-HD * Add comment about padding/masking * Fix checkpoint/repo references * Apply style fixes --------- Co-authored-by: github-actions[bot] Co-authored-by: Dhruv Nair --- .../en/api/models/chroma_transformer.md | 2 +- docs/source/en/api/pipelines/chroma.md | 13 +++++----- .../models/transformers/transformer_chroma.py | 2 +- .../pipelines/chroma/pipeline_chroma.py | 25 +++++++++++-------- .../chroma/pipeline_chroma_img2img.py | 23 +++++++++-------- 5 files changed, 35 insertions(+), 30 deletions(-) diff --git a/docs/source/en/api/models/chroma_transformer.md b/docs/source/en/api/models/chroma_transformer.md index 681e81f7a5..1ef24cda39 100644 --- a/docs/source/en/api/models/chroma_transformer.md +++ b/docs/source/en/api/models/chroma_transformer.md @@ -12,7 +12,7 @@ specific language governing permissions and limitations under the License. # ChromaTransformer2DModel -A modified flux Transformer model from [Chroma](https://huggingface.co/lodestones/Chroma) +A modified flux Transformer model from [Chroma](https://huggingface.co/lodestones/Chroma1-HD) ## ChromaTransformer2DModel diff --git a/docs/source/en/api/pipelines/chroma.md b/docs/source/en/api/pipelines/chroma.md index df03fbb325..cc52ffa09a 100644 --- a/docs/source/en/api/pipelines/chroma.md +++ b/docs/source/en/api/pipelines/chroma.md @@ -19,20 +19,21 @@ specific language governing permissions and limitations under the License. Chroma is a text to image generation model based on Flux. -Original model checkpoints for Chroma can be found [here](https://huggingface.co/lodestones/Chroma). +Original model checkpoints for Chroma can be found here: +* High-resolution finetune: [lodestones/Chroma1-HD](https://huggingface.co/lodestones/Chroma1-HD) +* Base model: [lodestones/Chroma1-Base](https://huggingface.co/lodestones/Chroma1-Base) +* Original repo with progress checkpoints: [lodestones/Chroma](https://huggingface.co/lodestones/Chroma) (loading this repo with `from_pretrained` will load a Diffusers-compatible version of the `unlocked-v37` checkpoint) > [!TIP] > Chroma can use all the same optimizations as Flux. ## Inference -The Diffusers version of Chroma is based on the [`unlocked-v37`](https://huggingface.co/lodestones/Chroma/blob/main/chroma-unlocked-v37.safetensors) version of the original model, which is available in the [Chroma repository](https://huggingface.co/lodestones/Chroma). - ```python import torch from diffusers import ChromaPipeline -pipe = ChromaPipeline.from_pretrained("lodestones/Chroma", torch_dtype=torch.bfloat16) +pipe = ChromaPipeline.from_pretrained("lodestones/Chroma1-HD", torch_dtype=torch.bfloat16) pipe.enable_model_cpu_offload() prompt = [ @@ -63,10 +64,10 @@ Then run the following example import torch from diffusers import ChromaTransformer2DModel, ChromaPipeline -model_id = "lodestones/Chroma" +model_id = "lodestones/Chroma1-HD" dtype = torch.bfloat16 -transformer = ChromaTransformer2DModel.from_single_file("https://huggingface.co/lodestones/Chroma/blob/main/chroma-unlocked-v37.safetensors", torch_dtype=dtype) +transformer = ChromaTransformer2DModel.from_single_file("https://huggingface.co/lodestones/Chroma1-HD/blob/main/Chroma1-HD.safetensors", torch_dtype=dtype) pipe = ChromaPipeline.from_pretrained(model_id, transformer=transformer, torch_dtype=dtype) pipe.enable_model_cpu_offload() diff --git a/src/diffusers/models/transformers/transformer_chroma.py b/src/diffusers/models/transformers/transformer_chroma.py index 5823ae9d3d..2ef3643daf 100644 --- a/src/diffusers/models/transformers/transformer_chroma.py +++ b/src/diffusers/models/transformers/transformer_chroma.py @@ -379,7 +379,7 @@ class ChromaTransformer2DModel( """ The Transformer model introduced in Flux, modified for Chroma. - Reference: https://huggingface.co/lodestones/Chroma + Reference: https://huggingface.co/lodestones/Chroma1-HD Args: patch_size (`int`, defaults to `1`): diff --git a/src/diffusers/pipelines/chroma/pipeline_chroma.py b/src/diffusers/pipelines/chroma/pipeline_chroma.py index 5482035b3a..ed6c2c2105 100644 --- a/src/diffusers/pipelines/chroma/pipeline_chroma.py +++ b/src/diffusers/pipelines/chroma/pipeline_chroma.py @@ -53,8 +53,8 @@ EXAMPLE_DOC_STRING = """ >>> import torch >>> from diffusers import ChromaPipeline - >>> model_id = "lodestones/Chroma" - >>> ckpt_path = "https://huggingface.co/lodestones/Chroma/blob/main/chroma-unlocked-v37.safetensors" + >>> model_id = "lodestones/Chroma1-HD" + >>> ckpt_path = "https://huggingface.co/lodestones/Chroma1-HD/blob/main/Chroma1-HD.safetensors" >>> transformer = ChromaTransformer2DModel.from_single_file(ckpt_path, torch_dtype=torch.bfloat16) >>> pipe = ChromaPipeline.from_pretrained( ... model_id, @@ -158,7 +158,7 @@ class ChromaPipeline( r""" The Chroma pipeline for text-to-image generation. - Reference: https://huggingface.co/lodestones/Chroma/ + Reference: https://huggingface.co/lodestones/Chroma1-HD/ Args: transformer ([`ChromaTransformer2DModel`]): @@ -233,20 +233,23 @@ class ChromaPipeline( return_tensors="pt", ) text_input_ids = text_inputs.input_ids - attention_mask = text_inputs.attention_mask.clone() + tokenizer_mask = text_inputs.attention_mask - # Chroma requires the attention mask to include one padding token - seq_lengths = attention_mask.sum(dim=1) - mask_indices = torch.arange(attention_mask.size(1)).unsqueeze(0).expand(batch_size, -1) - attention_mask = (mask_indices <= seq_lengths.unsqueeze(1)).bool() + tokenizer_mask_device = tokenizer_mask.to(device) + # unlike FLUX, Chroma uses the attention mask when generating the T5 embedding prompt_embeds = self.text_encoder( - text_input_ids.to(device), output_hidden_states=False, attention_mask=attention_mask.to(device) + text_input_ids.to(device), + output_hidden_states=False, + attention_mask=tokenizer_mask_device, )[0] - dtype = self.text_encoder.dtype prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) - attention_mask = attention_mask.to(device=device) + + # for the text tokens, chroma requires that all except the first padding token are masked out during the forward pass through the transformer + seq_lengths = tokenizer_mask_device.sum(dim=1) + mask_indices = torch.arange(tokenizer_mask_device.size(1), device=device).unsqueeze(0).expand(batch_size, -1) + attention_mask = (mask_indices <= seq_lengths.unsqueeze(1)).to(dtype=dtype, device=device) _, seq_len, _ = prompt_embeds.shape diff --git a/src/diffusers/pipelines/chroma/pipeline_chroma_img2img.py b/src/diffusers/pipelines/chroma/pipeline_chroma_img2img.py index 9afd4b9e15..470c746e41 100644 --- a/src/diffusers/pipelines/chroma/pipeline_chroma_img2img.py +++ b/src/diffusers/pipelines/chroma/pipeline_chroma_img2img.py @@ -53,8 +53,8 @@ EXAMPLE_DOC_STRING = """ >>> import torch >>> from diffusers import ChromaTransformer2DModel, ChromaImg2ImgPipeline - >>> model_id = "lodestones/Chroma" - >>> ckpt_path = "https://huggingface.co/lodestones/Chroma/blob/main/chroma-unlocked-v37.safetensors" + >>> model_id = "lodestones/Chroma1-HD" + >>> ckpt_path = "https://huggingface.co/lodestones/Chroma1-HD/blob/main/Chroma1-HD.safetensors" >>> pipe = ChromaImg2ImgPipeline.from_pretrained( ... model_id, ... transformer=transformer, @@ -170,7 +170,7 @@ class ChromaImg2ImgPipeline( r""" The Chroma pipeline for image-to-image generation. - Reference: https://huggingface.co/lodestones/Chroma/ + Reference: https://huggingface.co/lodestones/Chroma1-HD/ Args: transformer ([`ChromaTransformer2DModel`]): @@ -247,20 +247,21 @@ class ChromaImg2ImgPipeline( return_tensors="pt", ) text_input_ids = text_inputs.input_ids - attention_mask = text_inputs.attention_mask.clone() + tokenizer_mask = text_inputs.attention_mask - # Chroma requires the attention mask to include one padding token - seq_lengths = attention_mask.sum(dim=1) - mask_indices = torch.arange(attention_mask.size(1)).unsqueeze(0).expand(batch_size, -1) - attention_mask = (mask_indices <= seq_lengths.unsqueeze(1)).long() + tokenizer_mask_device = tokenizer_mask.to(device) prompt_embeds = self.text_encoder( - text_input_ids.to(device), output_hidden_states=False, attention_mask=attention_mask.to(device) + text_input_ids.to(device), + output_hidden_states=False, + attention_mask=tokenizer_mask_device, )[0] - dtype = self.text_encoder.dtype prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) - attention_mask = attention_mask.to(dtype=dtype, device=device) + + seq_lengths = tokenizer_mask_device.sum(dim=1) + mask_indices = torch.arange(tokenizer_mask_device.size(1), device=device).unsqueeze(0).expand(batch_size, -1) + attention_mask = (mask_indices <= seq_lengths.unsqueeze(1)).to(dtype=dtype, device=device) _, seq_len, _ = prompt_embeds.shape From 250f5cb53db1554f32dee07ad002f6c3834306d0 Mon Sep 17 00:00:00 2001 From: Mikko Lauri <167742591+lauri9@users.noreply.github.com> Date: Mon, 27 Oct 2025 16:55:02 +0200 Subject: [PATCH 097/123] Add AITER attention backend (#12549) * add aiter attention backend * Apply style fixes --------- Co-authored-by: Sayak Paul Co-authored-by: github-actions[bot] --- .../en/optimization/attention_backends.md | 2 + src/diffusers/models/attention_dispatch.py | 60 +++++++++++++++++++ src/diffusers/utils/__init__.py | 2 + src/diffusers/utils/import_utils.py | 21 +++++++ tests/others/test_attention_backends.py | 13 ++++ 5 files changed, 98 insertions(+) diff --git a/docs/source/en/optimization/attention_backends.md b/docs/source/en/optimization/attention_backends.md index 8be2c06030..edfdcc38b5 100644 --- a/docs/source/en/optimization/attention_backends.md +++ b/docs/source/en/optimization/attention_backends.md @@ -21,6 +21,7 @@ Refer to the table below for an overview of the available attention families and | attention family | main feature | |---|---| | FlashAttention | minimizes memory reads/writes through tiling and recomputation | +| AI Tensor Engine for ROCm | FlashAttention implementation optimized for AMD ROCm accelerators | | SageAttention | quantizes attention to int8 | | PyTorch native | built-in PyTorch implementation using [scaled_dot_product_attention](./fp16#scaled-dot-product-attention) | | xFormers | memory-efficient attention with support for various attention kernels | @@ -139,6 +140,7 @@ Refer to the table below for a complete list of available attention backends and | `_native_xla` | [PyTorch native](https://docs.pytorch.org/docs/stable/generated/torch.nn.attention.SDPBackend.html#torch.nn.attention.SDPBackend) | XLA-optimized attention | | `flash` | [FlashAttention](https://github.com/Dao-AILab/flash-attention) | FlashAttention-2 | | `flash_varlen` | [FlashAttention](https://github.com/Dao-AILab/flash-attention) | Variable length FlashAttention | +| `aiter` | [AI Tensor Engine for ROCm](https://github.com/ROCm/aiter) | FlashAttention for AMD ROCm | | `_flash_3` | [FlashAttention](https://github.com/Dao-AILab/flash-attention) | FlashAttention-3 | | `_flash_varlen_3` | [FlashAttention](https://github.com/Dao-AILab/flash-attention) | Variable length FlashAttention-3 | | `_flash_3_hub` | [FlashAttention](https://github.com/Dao-AILab/flash-attention) | FlashAttention-3 from kernels | diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index e169491099..ab0d7102ee 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -27,6 +27,8 @@ if torch.distributed.is_available(): from ..utils import ( get_logger, + is_aiter_available, + is_aiter_version, is_flash_attn_3_available, is_flash_attn_available, is_flash_attn_version, @@ -47,6 +49,7 @@ if TYPE_CHECKING: from ._modeling_parallel import ParallelConfig _REQUIRED_FLASH_VERSION = "2.6.3" +_REQUIRED_AITER_VERSION = "0.1.5" _REQUIRED_SAGE_VERSION = "2.1.1" _REQUIRED_FLEX_VERSION = "2.5.0" _REQUIRED_XLA_VERSION = "2.2" @@ -54,6 +57,7 @@ _REQUIRED_XFORMERS_VERSION = "0.0.29" _CAN_USE_FLASH_ATTN = is_flash_attn_available() and is_flash_attn_version(">=", _REQUIRED_FLASH_VERSION) _CAN_USE_FLASH_ATTN_3 = is_flash_attn_3_available() +_CAN_USE_AITER_ATTN = is_aiter_available() and is_aiter_version(">=", _REQUIRED_AITER_VERSION) _CAN_USE_SAGE_ATTN = is_sageattention_available() and is_sageattention_version(">=", _REQUIRED_SAGE_VERSION) _CAN_USE_FLEX_ATTN = is_torch_version(">=", _REQUIRED_FLEX_VERSION) _CAN_USE_NPU_ATTN = is_torch_npu_available() @@ -78,6 +82,12 @@ else: flash_attn_3_func = None flash_attn_3_varlen_func = None + +if _CAN_USE_AITER_ATTN: + from aiter import flash_attn_func as aiter_flash_attn_func +else: + aiter_flash_attn_func = None + if DIFFUSERS_ENABLE_HUB_KERNELS: if not is_kernels_available(): raise ImportError( @@ -178,6 +188,9 @@ class AttentionBackendName(str, Enum): _FLASH_3_HUB = "_flash_3_hub" # _FLASH_VARLEN_3_HUB = "_flash_varlen_3_hub" # not supported yet. + # `aiter` + AITER = "aiter" + # PyTorch native FLEX = "flex" NATIVE = "native" @@ -414,6 +427,12 @@ def _check_attention_backend_requirements(backend: AttentionBackendName) -> None f"Flash Attention 3 Hub backend '{backend.value}' is not usable because the `kernels` package isn't available. Please install it with `pip install kernels`." ) + elif backend == AttentionBackendName.AITER: + if not _CAN_USE_AITER_ATTN: + raise RuntimeError( + f"Aiter Attention backend '{backend.value}' is not usable because of missing package or the version is too old. Please install `aiter>={_REQUIRED_AITER_VERSION}`." + ) + elif backend in [ AttentionBackendName.SAGE, AttentionBackendName.SAGE_VARLEN, @@ -1397,6 +1416,47 @@ def _flash_varlen_attention_3( return (out, lse) if return_lse else out +@_AttentionBackendRegistry.register( + AttentionBackendName.AITER, + constraints=[_check_device_cuda, _check_qkv_dtype_bf16_or_fp16, _check_shape], +) +def _aiter_flash_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: Optional[float] = None, + return_lse: bool = False, + _parallel_config: Optional["ParallelConfig"] = None, +) -> torch.Tensor: + if not return_lse and torch.is_grad_enabled(): + # aiter requires return_lse=True by assertion when gradients are enabled. + out, lse, *_ = aiter_flash_attn_func( + q=query, + k=key, + v=value, + dropout_p=dropout_p, + softmax_scale=scale, + causal=is_causal, + return_lse=True, + ) + else: + out = aiter_flash_attn_func( + q=query, + k=key, + v=value, + dropout_p=dropout_p, + softmax_scale=scale, + causal=is_causal, + return_lse=return_lse, + ) + if return_lse: + out, lse, *_ = out + + return (out, lse) if return_lse else out + + @_AttentionBackendRegistry.register( AttentionBackendName.FLEX, constraints=[_check_attn_mask_or_causal, _check_device, _check_shape], diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py index d8e1a55401..cf77aaee82 100644 --- a/src/diffusers/utils/__init__.py +++ b/src/diffusers/utils/__init__.py @@ -64,6 +64,8 @@ from .import_utils import ( get_objects_from_module, is_accelerate_available, is_accelerate_version, + is_aiter_available, + is_aiter_version, is_better_profanity_available, is_bitsandbytes_available, is_bitsandbytes_version, diff --git a/src/diffusers/utils/import_utils.py b/src/diffusers/utils/import_utils.py index 97065267b0..adf8ed8b06 100644 --- a/src/diffusers/utils/import_utils.py +++ b/src/diffusers/utils/import_utils.py @@ -226,6 +226,7 @@ _cosmos_guardrail_available, _cosmos_guardrail_version = _is_package_available(" _sageattention_available, _sageattention_version = _is_package_available("sageattention") _flash_attn_available, _flash_attn_version = _is_package_available("flash_attn") _flash_attn_3_available, _flash_attn_3_version = _is_package_available("flash_attn_3") +_aiter_available, _aiter_version = _is_package_available("aiter") _kornia_available, _kornia_version = _is_package_available("kornia") _nvidia_modelopt_available, _nvidia_modelopt_version = _is_package_available("modelopt", get_dist_name=True) @@ -406,6 +407,10 @@ def is_flash_attn_3_available(): return _flash_attn_3_available +def is_aiter_available(): + return _aiter_available + + def is_kornia_available(): return _kornia_available @@ -911,6 +916,22 @@ def is_flash_attn_version(operation: str, version: str): return compare_versions(parse(_flash_attn_version), operation, version) +@cache +def is_aiter_version(operation: str, version: str): + """ + Compares the current aiter version to a given reference with an operation. + + Args: + operation (`str`): + A string representation of an operator, such as `">"` or `"<="` + version (`str`): + A version string + """ + if not _aiter_available: + return False + return compare_versions(parse(_aiter_version), operation, version) + + def get_objects_from_module(module): """ Returns a dict of object names and values in a module, while skipping private/internal objects diff --git a/tests/others/test_attention_backends.py b/tests/others/test_attention_backends.py index 42cdcd56f7..2e5a2fc82b 100644 --- a/tests/others/test_attention_backends.py +++ b/tests/others/test_attention_backends.py @@ -14,6 +14,10 @@ pytest tests/others/test_attention_backends.py Tests were conducted on an H100 with PyTorch 2.8.0 (CUDA 12.9). Slices for the compilation tests in "native" variants were obtained with a torch nightly version (2.10.0.dev20250924+cu128). + +Tests for aiter backend were conducted and slices for the aiter backend tests collected on a MI355X +with torch 2025-09-25 nightly version (ad2f7315ca66b42497047bb7951f696b50f1e81b) and +aiter 0.1.5.post4.dev20+ga25e55e79. """ import os @@ -44,6 +48,10 @@ FORWARD_CASES = [ "_native_cudnn", torch.tensor([0.0781, 0.0840, 0.0879, 0.0957, 0.0898, 0.0957, 0.0957, 0.0977, 0.2168, 0.2246, 0.2324, 0.2500, 0.2539, 0.2480, 0.2441, 0.2695], dtype=torch.bfloat16), ), + ( + "aiter", + torch.tensor([0.0781, 0.0820, 0.0879, 0.0957, 0.0898, 0.0938, 0.0957, 0.0957, 0.2285, 0.2363, 0.2461, 0.2637, 0.2695, 0.2617, 0.2617, 0.2891], dtype=torch.bfloat16), + ) ] COMPILE_CASES = [ @@ -63,6 +71,11 @@ COMPILE_CASES = [ torch.tensor([0.0410, 0.0410, 0.0430, 0.0508, 0.0488, 0.0586, 0.0605, 0.0586, 0.2344, 0.2461, 0.2578, 0.2773, 0.2871, 0.2832, 0.2793, 0.3086], dtype=torch.bfloat16), True, ), + ( + "aiter", + torch.tensor([0.0391, 0.0391, 0.0430, 0.0488, 0.0469, 0.0566, 0.0586, 0.0566, 0.2402, 0.2539, 0.2637, 0.2812, 0.2930, 0.2910, 0.2891, 0.3164], dtype=torch.bfloat16), + True, + ) ] # fmt: on From 6d1a6486024192951ce696e8f4cf79a39509182f Mon Sep 17 00:00:00 2001 From: alirezafarashah <63104907+alirezafarashah@users.noreply.github.com> Date: Mon, 27 Oct 2025 13:16:43 -0400 Subject: [PATCH 098/123] Fix small inconsistency in output dimension of "_get_t5_prompt_embeds" function in sd3 pipeline (#12531) * Fix small inconsistency in output dimension of t5 embeds when text_encoder_3 is None * first commit --------- Co-authored-by: Alireza Farashah Co-authored-by: Alireza Farashah --- .../controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py | 2 +- .../pipeline_stable_diffusion_3_controlnet_inpainting.py | 2 +- src/diffusers/pipelines/pag/pipeline_pag_sd_3.py | 2 +- src/diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py | 2 +- .../pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py | 2 +- .../stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py | 2 +- .../stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py | 2 +- 7 files changed, 7 insertions(+), 7 deletions(-) diff --git a/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py b/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py index c763411ab5..f67a0e2112 100644 --- a/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py +++ b/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py @@ -266,7 +266,7 @@ class StableDiffusion3ControlNetPipeline( return torch.zeros( ( batch_size * num_images_per_prompt, - self.tokenizer_max_length, + max_sequence_length, self.transformer.config.joint_attention_dim, ), device=device, diff --git a/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py b/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py index c33cf979c6..68984da4dc 100644 --- a/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py +++ b/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py @@ -284,7 +284,7 @@ class StableDiffusion3ControlNetInpaintingPipeline( return torch.zeros( ( batch_size * num_images_per_prompt, - self.tokenizer_max_length, + max_sequence_length, self.transformer.config.joint_attention_dim, ), device=device, diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py index acb4e52340..bc281428e2 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py @@ -237,7 +237,7 @@ class StableDiffusion3PAGPipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSin return torch.zeros( ( batch_size * num_images_per_prompt, - self.tokenizer_max_length, + max_sequence_length, self.transformer.config.joint_attention_dim, ), device=device, diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py index e1819a79fb..22a8dac238 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py @@ -253,7 +253,7 @@ class StableDiffusion3PAGImg2ImgPipeline(DiffusionPipeline, SD3LoraLoaderMixin, return torch.zeros( ( batch_size * num_images_per_prompt, - self.tokenizer_max_length, + max_sequence_length, self.transformer.config.joint_attention_dim, ), device=device, diff --git a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py index 1618f89a49..3b7b26dc63 100644 --- a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py +++ b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py @@ -248,7 +248,7 @@ class StableDiffusion3Pipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingle return torch.zeros( ( batch_size * num_images_per_prompt, - self.tokenizer_max_length, + max_sequence_length, self.transformer.config.joint_attention_dim, ), device=device, diff --git a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py index 7e97909f42..db047f1992 100644 --- a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py +++ b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py @@ -272,7 +272,7 @@ class StableDiffusion3Img2ImgPipeline(DiffusionPipeline, SD3LoraLoaderMixin, Fro return torch.zeros( ( batch_size * num_images_per_prompt, - self.tokenizer_max_length, + max_sequence_length, self.transformer.config.joint_attention_dim, ), device=device, diff --git a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py index 5b2cca0378..c95fa530c8 100644 --- a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py +++ b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py @@ -278,7 +278,7 @@ class StableDiffusion3InpaintPipeline(DiffusionPipeline, SD3LoraLoaderMixin, Fro return torch.zeros( ( batch_size * num_images_per_prompt, - self.tokenizer_max_length, + max_sequence_length, self.transformer.config.joint_attention_dim, ), device=device, From 5afbcce176cd4e8ec08f43ee9fae2d6562edf54c Mon Sep 17 00:00:00 2001 From: Lev Novitskiy <57654885+leffff@users.noreply.github.com> Date: Tue, 28 Oct 2025 05:17:18 +0300 Subject: [PATCH 099/123] Kandinsky 5 10 sec (NABLA suport) (#12520) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * add transformer pipeline first version * updates * fix 5sec generation * rewrite Kandinsky5T2VPipeline to diffusers style * add multiprompt support * remove prints in pipeline * add nabla attention * Wrap Transformer in Diffusers style * fix license * fix prompt type * add gradient checkpointing and peft support * add usage example * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: Álvaro Somoza * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: Álvaro Somoza * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: Álvaro Somoza * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: Álvaro Somoza * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: Álvaro Somoza * remove unused imports * add 10 second models support * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * remove no_grad and simplified prompt paddings * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * moved template to __init__ * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: YiYi Xu * moved sdps inside processor * remove oneline function * remove reset_dtype methods * Transformer: move all methods to forward * separated prompt encoding * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: YiYi Xu * refactoring * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: YiYi Xu * refactoring acording to https://github.com/huggingface/diffusers/commit/acabbc0033d4b4933fc651766a4aa026db2e6dc1 * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * fixed * style +copies * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: Charles * more * Apply suggestions from code review * add lora loader doc * add compiled Nabla Attention * all needed changes for 10 sec models are added! * add docs * Apply style fixes * update docs * add kandinsky5 to toctree * add tests * fix tests * Apply style fixes * update tests --------- Co-authored-by: Álvaro Somoza Co-authored-by: YiYi Xu Co-authored-by: Charles Co-authored-by: Sayak Paul Co-authored-by: github-actions[bot] --- docs/source/en/_toctree.yml | 2 + docs/source/en/api/pipelines/kandinsky5.md | 149 +++++++++ .../transformers/transformer_kandinsky.py | 2 + .../kandinsky5/pipeline_kandinsky.py | 9 +- tests/pipelines/kandinsky5/__init__.py | 0 tests/pipelines/kandinsky5/test_kandinsky5.py | 306 ++++++++++++++++++ tests/pipelines/test_pipelines_common.py | 2 + 7 files changed, 468 insertions(+), 2 deletions(-) create mode 100644 docs/source/en/api/pipelines/kandinsky5.md create mode 100644 tests/pipelines/kandinsky5/__init__.py create mode 100644 tests/pipelines/kandinsky5/test_kandinsky5.py diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 540e99a2c6..44870f680e 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -525,6 +525,8 @@ title: Kandinsky 2.2 - local: api/pipelines/kandinsky3 title: Kandinsky 3 + - local: api/pipelines/kandinsky5 + title: Kandinsky 5 - local: api/pipelines/kolors title: Kolors - local: api/pipelines/latent_consistency_models diff --git a/docs/source/en/api/pipelines/kandinsky5.md b/docs/source/en/api/pipelines/kandinsky5.md new file mode 100644 index 0000000000..cb1c119f80 --- /dev/null +++ b/docs/source/en/api/pipelines/kandinsky5.md @@ -0,0 +1,149 @@ + + +# Kandinsky 5.0 + +Kandinsky 5.0 is created by the Kandinsky team: Alexey Letunovskiy, Maria Kovaleva, Ivan Kirillov, Lev Novitskiy, Denis Koposov, Dmitrii Mikhailov, Anna Averchenkova, Andrey Shutkin, Julia Agafonova, Olga Kim, Anastasiia Kargapoltseva, Nikita Kiselev, Anna Dmitrienko, Anastasia Maltseva, Kirill Chernyshev, Ilia Vasiliev, Viacheslav Vasilev, Vladimir Polovnikov, Yury Kolabushin, Alexander Belykh, Mikhail Mamaev, Anastasia Aliaskina, Tatiana Nikulina, Polina Gavrilova, Vladimir Arkhipkin, Vladimir Korviakov, Nikolai Gerasimenko, Denis Parkhomenko, Denis Dimitrov + + +Kandinsky 5.0 is a family of diffusion models for Video & Image generation. Kandinsky 5.0 T2V Lite is a lightweight video generation model (2B parameters) that ranks #1 among open-source models in its class. It outperforms larger models and offers the best understanding of Russian concepts in the open-source ecosystem. + +The model introduces several key innovations: +- **Latent diffusion pipeline** with **Flow Matching** for improved training stability +- **Diffusion Transformer (DiT)** as the main generative backbone with cross-attention to text embeddings +- Dual text encoding using **Qwen2.5-VL** and **CLIP** for comprehensive text understanding +- **HunyuanVideo 3D VAE** for efficient video encoding and decoding +- **Sparse attention mechanisms** (NABLA) for efficient long-sequence processing + +The original codebase can be found at [ai-forever/Kandinsky-5](https://github.com/ai-forever/Kandinsky-5). + +> [!TIP] +> Check out the [AI Forever](https://huggingface.co/ai-forever) organization on the Hub for the official model checkpoints for text-to-video generation, including pretrained, SFT, no-CFG, and distilled variants. + +## Available Models + +Kandinsky 5.0 T2V Lite comes in several variants optimized for different use cases: + +| model_id | Description | Use Cases | +|------------|-------------|-----------| +| **ai-forever/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers** | 5 second Supervised Fine-Tuned model | Highest generation quality | +| **ai-forever/Kandinsky-5.0-T2V-Lite-sft-10s-Diffusers** | 10 second Supervised Fine-Tuned model | Highest generation quality | +| **ai-forever/Kandinsky-5.0-T2V-Lite-nocfg-5s-Diffusers** | 5 second Classifier-Free Guidance distilled | 2× faster inference | +| **ai-forever/Kandinsky-5.0-T2V-Lite-nocfg-10s-Diffusers** | 10 second Classifier-Free Guidance distilled | 2× faster inference | +| **ai-forever/Kandinsky-5.0-T2V-Lite-distilled16steps-5s-Diffusers** | 5 second Diffusion distilled to 16 steps | 6× faster inference, minimal quality loss | +| **ai-forever/Kandinsky-5.0-T2V-Lite-distilled16steps-10s-Diffusers** | 10 second Diffusion distilled to 16 steps | 6× faster inference, minimal quality loss | +| **ai-forever/Kandinsky-5.0-T2V-Lite-pretrain-5s-Diffusers** | 5 second Base pretrained model | Research and fine-tuning | +| **ai-forever/Kandinsky-5.0-T2V-Lite-pretrain-10s-Diffusers** | 10 second Base pretrained model | Research and fine-tuning | + +All models are available in 5-second and 10-second video generation versions. + +## Kandinsky5T2VPipeline + +[[autodoc]] Kandinsky5T2VPipeline + - all + - __call__ + +## Usage Examples + +### Basic Text-to-Video Generation + +```python +import torch +from diffusers import Kandinsky5T2VPipeline +from diffusers.utils import export_to_video + +# Load the pipeline +model_id = "ai-forever/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers" +pipe = Kandinsky5T2VPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16) +pipe = pipe.to("cuda") + +# Generate video +prompt = "A cat and a dog baking a cake together in a kitchen." +negative_prompt = "Static, 2D cartoon, cartoon, 2d animation, paintings, images, worst quality, low quality, ugly, deformed, walking backwards" + +output = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + height=512, + width=768, + num_frames=121, # ~5 seconds at 24fps + num_inference_steps=50, + guidance_scale=5.0, +).frames[0] + +export_to_video(output, "output.mp4", fps=24, quality=9) +``` + +### 10 second Models +**⚠️ Warning!** all 10 second models should be used with Flex attention and max-autotune-no-cudagraphs compilation: + +```python +pipe = Kandinsky5T2VPipeline.from_pretrained( + "ai-forever/Kandinsky-5.0-T2V-Lite-sft-10s-Diffusers", + torch_dtype=torch.bfloat16 +) +pipe = pipe.to("cuda") + +pipe.transformer.set_attention_backend( + "flex" +) # <--- Sett attention bakend to Flex +pipe.transformer.compile( + mode="max-autotune-no-cudagraphs", + dynamic=True +) # <--- Compile with max-autotune-no-cudagraphs + +prompt = "A cat and a dog baking a cake together in a kitchen." +negative_prompt = "Static, 2D cartoon, cartoon, 2d animation, paintings, images, worst quality, low quality, ugly, deformed, walking backwards" + +output = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + height=512, + width=768, + num_frames=241, + num_inference_steps=50, + guidance_scale=5.0, +).frames[0] + +export_to_video(output, "output.mp4", fps=24, quality=9) +``` + +### Diffusion Distilled model +**⚠️ Warning!** all nocfg and diffusion distilled models should be infered wothout CFG (```guidance_scale=1.0```): + +```python +model_id = "ai-forever/Kandinsky-5.0-T2V-Lite-distilled16steps-5s-Diffusers" +pipe = Kandinsky5T2VPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16) +pipe = pipe.to("cuda") + +output = pipe( + prompt="A beautiful sunset over mountains", + num_inference_steps=16, # <--- Model is distilled in 16 steps + guidance_scale=1.0, # <--- no CFG +).frames[0] + +export_to_video(output, "output.mp4", fps=24, quality=9) +``` + + +## Citation +```bibtex +@misc{kandinsky2025, + author = {Alexey Letunovskiy and Maria Kovaleva and Ivan Kirillov and Lev Novitskiy and Denis Koposov and + Dmitrii Mikhailov and Anna Averchenkova and Andrey Shutkin and Julia Agafonova and Olga Kim and + Anastasiia Kargapoltseva and Nikita Kiselev and Vladimir Arkhipkin and Vladimir Korviakov and + Nikolai Gerasimenko and Denis Parkhomenko and Anna Dmitrienko and Anastasia Maltseva and + Kirill Chernyshev and Ilia Vasiliev and Viacheslav Vasilev and Vladimir Polovnikov and + Yury Kolabushin and Alexander Belykh and Mikhail Mamaev and Anastasia Aliaskina and + Tatiana Nikulina and Polina Gavrilova and Denis Dimitrov}, + title = {Kandinsky 5.0: A family of diffusion models for Video & Image generation}, + howpublished = {\url{https://github.com/ai-forever/Kandinsky-5}}, + year = 2025 +} +``` \ No newline at end of file diff --git a/src/diffusers/models/transformers/transformer_kandinsky.py b/src/diffusers/models/transformers/transformer_kandinsky.py index d4ba92acaf..316e79da4f 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky.py +++ b/src/diffusers/models/transformers/transformer_kandinsky.py @@ -324,6 +324,7 @@ class Kandinsky5AttnProcessor: sparse_params["sta_mask"], thr=sparse_params["P"], ) + else: attn_mask = None @@ -335,6 +336,7 @@ class Kandinsky5AttnProcessor: backend=self._attention_backend, parallel_config=self._parallel_config, ) + hidden_states = hidden_states.flatten(-2, -1) attn_out = attn.out_layer(hidden_states) diff --git a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py index 2b977a5a36..3f93aa1889 100644 --- a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py +++ b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py @@ -173,8 +173,10 @@ class Kandinsky5T2VPipeline(DiffusionPipeline, KandinskyLoraLoaderMixin): ) self.prompt_template_encode_start_idx = 129 - self.vae_scale_factor_temporal = vae.config.temporal_compression_ratio - self.vae_scale_factor_spatial = vae.config.spatial_compression_ratio + self.vae_scale_factor_temporal = ( + self.vae.config.temporal_compression_ratio if getattr(self, "vae", None) else 4 + ) + self.vae_scale_factor_spatial = self.vae.config.spatial_compression_ratio if getattr(self, "vae", None) else 8 self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) @staticmethod @@ -384,6 +386,9 @@ class Kandinsky5T2VPipeline(DiffusionPipeline, KandinskyLoraLoaderMixin): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype + if not isinstance(prompt, list): + prompt = [prompt] + batch_size = len(prompt) prompt = [prompt_clean(p) for p in prompt] diff --git a/tests/pipelines/kandinsky5/__init__.py b/tests/pipelines/kandinsky5/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/pipelines/kandinsky5/test_kandinsky5.py b/tests/pipelines/kandinsky5/test_kandinsky5.py new file mode 100644 index 0000000000..47fccb632a --- /dev/null +++ b/tests/pipelines/kandinsky5/test_kandinsky5.py @@ -0,0 +1,306 @@ +# Copyright 2025 The Kandinsky Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import torch +from transformers import ( + CLIPTextConfig, + CLIPTextModel, + CLIPTokenizer, + Qwen2_5_VLConfig, + Qwen2_5_VLForConditionalGeneration, + Qwen2VLProcessor, +) + +from diffusers import ( + AutoencoderKLHunyuanVideo, + FlowMatchEulerDiscreteScheduler, + Kandinsky5T2VPipeline, + Kandinsky5Transformer3DModel, +) + +from ...testing_utils import ( + enable_full_determinism, + torch_device, +) +from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS +from ..test_pipelines_common import PipelineTesterMixin + + +enable_full_determinism() + + +class Kandinsky5T2VPipelineFastTests(PipelineTesterMixin, unittest.TestCase): + pipeline_class = Kandinsky5T2VPipeline + params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs", "prompt_embeds", "negative_prompt_embeds"} + batch_params = TEXT_TO_IMAGE_BATCH_PARAMS + + # Define required optional parameters for your pipeline + required_optional_params = frozenset( + [ + "num_inference_steps", + "generator", + "latents", + "return_dict", + "callback_on_step_end", + "callback_on_step_end_tensor_inputs", + "max_sequence_length", + ] + ) + + test_xformers_attention = False + supports_dduf = False + + def get_dummy_components(self): + torch.manual_seed(0) + vae = AutoencoderKLHunyuanVideo( + in_channels=3, + out_channels=3, + spatial_compression_ratio=8, + temporal_compression_ratio=4, + latent_channels=4, + block_out_channels=(8, 8, 8, 8), + layers_per_block=1, + norm_num_groups=4, + ) + + torch.manual_seed(0) + scheduler = FlowMatchEulerDiscreteScheduler(shift=7.0) + + # Dummy Qwen2.5-VL model + config = Qwen2_5_VLConfig( + text_config={ + "hidden_size": 16, + "intermediate_size": 16, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 2, + "rope_scaling": { + "mrope_section": [1, 1, 2], + "rope_type": "default", + "type": "default", + }, + "rope_theta": 1000000.0, + }, + vision_config={ + "depth": 2, + "hidden_size": 16, + "intermediate_size": 16, + "num_heads": 2, + "out_hidden_size": 16, + }, + hidden_size=16, + vocab_size=152064, + vision_end_token_id=151653, + vision_start_token_id=151652, + vision_token_id=151654, + ) + text_encoder = Qwen2_5_VLForConditionalGeneration(config) + tokenizer = Qwen2VLProcessor.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration") + + # Dummy CLIP model + clip_text_encoder_config = CLIPTextConfig( + bos_token_id=0, + eos_token_id=2, + hidden_size=32, + intermediate_size=37, + layer_norm_eps=1e-05, + num_attention_heads=4, + num_hidden_layers=5, + pad_token_id=1, + vocab_size=1000, + hidden_act="gelu", + projection_dim=32, + ) + + torch.manual_seed(0) + text_encoder_2 = CLIPTextModel(clip_text_encoder_config) + tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") + + torch.manual_seed(0) + transformer = Kandinsky5Transformer3DModel( + in_visual_dim=4, + in_text_dim=16, # Match tiny Qwen2.5-VL hidden size + in_text_dim2=32, # Match tiny CLIP hidden size + time_dim=32, + out_visual_dim=4, + patch_size=(1, 2, 2), + model_dim=48, + ff_dim=128, + num_text_blocks=1, + num_visual_blocks=1, + axes_dims=(8, 8, 8), + visual_cond=False, + ) + + components = { + "transformer": transformer.eval(), + "vae": vae.eval(), + "scheduler": scheduler, + "text_encoder": text_encoder.eval(), + "tokenizer": tokenizer, + "text_encoder_2": text_encoder_2.eval(), + "tokenizer_2": tokenizer_2, + } + return components + + def get_dummy_inputs(self, device, seed=0): + if str(device).startswith("mps"): + generator = torch.manual_seed(seed) + else: + generator = torch.Generator(device=device).manual_seed(seed) + inputs = { + "prompt": "A cat dancing", + "negative_prompt": "blurry, low quality", + "generator": generator, + "num_inference_steps": 2, + "guidance_scale": 5.0, + "height": 32, + "width": 32, + "num_frames": 5, + "max_sequence_length": 16, + "output_type": "pt", + } + return inputs + + def test_inference(self): + device = "cpu" + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + video = pipe(**inputs).frames + + # Check video shape: (batch, frames, channel, height, width) + expected_shape = (1, 5, 3, 32, 32) + self.assertEqual(video.shape, expected_shape) + + # Check specific values + expected_slice = torch.tensor( + [ + 0.4330, + 0.4254, + 0.4285, + 0.3835, + 0.4253, + 0.4196, + 0.3704, + 0.3714, + 0.4999, + 0.5346, + 0.4795, + 0.4637, + 0.4930, + 0.5124, + 0.4902, + 0.4570, + ] + ) + + generated_slice = video.flatten() + # Take first 8 and last 8 values for comparison + video_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) + self.assertTrue( + torch.allclose(video_slice, expected_slice, atol=1e-3), + f"video_slice: {video_slice}, expected_slice: {expected_slice}", + ) + + def test_inference_batch_single_identical(self): + # Override to test batch single identical with video + super().test_inference_batch_single_identical(batch_size=2, expected_max_diff=1e-2) + + def test_encode_prompt_works_in_isolation(self, extra_required_param_value_dict=None, atol=1e-3, rtol=1e-3): + components = self.get_dummy_components() + + text_component_names = ["text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2"] + text_components = {k: (v if k in text_component_names else None) for k, v in components.items()} + non_text_components = {k: (v if k not in text_component_names else None) for k, v in components.items()} + + pipe_with_just_text_encoder = self.pipeline_class(**text_components) + pipe_with_just_text_encoder = pipe_with_just_text_encoder.to(torch_device) + + pipe_without_text_encoders = self.pipeline_class(**non_text_components) + pipe_without_text_encoders = pipe_without_text_encoders.to(torch_device) + + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + + # Compute `encode_prompt()`. + + # Test single prompt + prompt = "A cat dancing" + with torch.no_grad(): + prompt_embeds_qwen, prompt_embeds_clip, prompt_cu_seqlens = pipe_with_just_text_encoder.encode_prompt( + prompt, device=torch_device, max_sequence_length=16 + ) + + # Check shapes + self.assertEqual(prompt_embeds_qwen.shape, (1, 4, 16)) # [batch, seq_len, embed_dim] + self.assertEqual(prompt_embeds_clip.shape, (1, 32)) # [batch, embed_dim] + self.assertEqual(prompt_cu_seqlens.shape, (2,)) # [batch + 1] + + # Test batch of prompts + prompts = ["A cat dancing", "A dog running"] + with torch.no_grad(): + batch_embeds_qwen, batch_embeds_clip, batch_cu_seqlens = pipe_with_just_text_encoder.encode_prompt( + prompts, device=torch_device, max_sequence_length=16 + ) + + # Check batch size + self.assertEqual(batch_embeds_qwen.shape, (len(prompts), 4, 16)) + self.assertEqual(batch_embeds_clip.shape, (len(prompts), 32)) + self.assertEqual(len(batch_cu_seqlens), len(prompts) + 1) # [0, len1, len1+len2] + + inputs = self.get_dummy_inputs(torch_device) + inputs["guidance_scale"] = 1.0 + + # baseline output: full pipeline + pipe_out = pipe(**inputs).frames + + # test against pipeline call with pre-computed prompt embeds + inputs = self.get_dummy_inputs(torch_device) + inputs["guidance_scale"] = 1.0 + + with torch.no_grad(): + prompt_embeds_qwen, prompt_embeds_clip, prompt_cu_seqlens = pipe_with_just_text_encoder.encode_prompt( + inputs["prompt"], device=torch_device, max_sequence_length=inputs["max_sequence_length"] + ) + + inputs["prompt"] = None + inputs["prompt_embeds_qwen"] = prompt_embeds_qwen + inputs["prompt_embeds_clip"] = prompt_embeds_clip + inputs["prompt_cu_seqlens"] = prompt_cu_seqlens + + pipe_out_2 = pipe_without_text_encoders(**inputs)[0] + + self.assertTrue( + torch.allclose(pipe_out, pipe_out_2, atol=atol, rtol=rtol), + f"max diff: {torch.max(torch.abs(pipe_out - pipe_out_2))}", + ) + + @unittest.skip("Kandinsky5T2VPipeline does not support attention slicing") + def test_attention_slicing_forward_pass(self): + pass + + @unittest.skip("Kandinsky5T2VPipeline does not support xformers") + def test_xformers_attention_forwardGenerator_pass(self): + pass + + @unittest.skip("Kandinsky5T2VPipeline does not support VAE slicing") + def test_vae_slicing(self): + pass diff --git a/tests/pipelines/test_pipelines_common.py b/tests/pipelines/test_pipelines_common.py index db8209835b..2af4ad0314 100644 --- a/tests/pipelines/test_pipelines_common.py +++ b/tests/pipelines/test_pipelines_common.py @@ -1461,6 +1461,8 @@ class PipelineTesterMixin: def test_save_load_optional_components(self, expected_max_difference=1e-4): if not hasattr(self.pipeline_class, "_optional_components"): return + if not self.pipeline_class._optional_components: + return components = self.get_dummy_components() pipe = self.pipeline_class(**components) for component in pipe.components.values(): From 303efd2b8d9949758004158c89b6c29a816150a2 Mon Sep 17 00:00:00 2001 From: "G.O.D" <32255912+gameofdimension@users.noreply.github.com> Date: Tue, 28 Oct 2025 10:55:36 +0800 Subject: [PATCH 100/123] Improve pos embed for Flux.1 inference on Ascend NPU (#12534) improve pos embed for ascend npu Co-authored-by: felix01.yu --- src/diffusers/models/transformers/transformer_flux.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_flux.py b/src/diffusers/models/transformers/transformer_flux.py index 1a44644324..16c526f437 100644 --- a/src/diffusers/models/transformers/transformer_flux.py +++ b/src/diffusers/models/transformers/transformer_flux.py @@ -22,7 +22,7 @@ import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FluxTransformer2DLoadersMixin, FromOriginalModelMixin, PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers +from ...utils import USE_PEFT_BACKEND, is_torch_npu_available, logging, scale_lora_layers, unscale_lora_layers from ...utils.torch_utils import maybe_allow_in_graph from .._modeling_parallel import ContextParallelInput, ContextParallelOutput from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward @@ -717,7 +717,11 @@ class FluxTransformer2DModel( img_ids = img_ids[0] ids = torch.cat((txt_ids, img_ids), dim=0) - image_rotary_emb = self.pos_embed(ids) + if is_torch_npu_available(): + freqs_cos, freqs_sin = self.pos_embed(ids.cpu()) + image_rotary_emb = (freqs_cos.npu(), freqs_sin.npu()) + else: + image_rotary_emb = self.pos_embed(ids) if joint_attention_kwargs is not None and "ip_adapter_image_embeds" in joint_attention_kwargs: ip_adapter_image_embeds = joint_attention_kwargs.pop("ip_adapter_image_embeds") From df0e2a4f2c2ebd7b93e6f51602d8d054afd8cd86 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Tue, 28 Oct 2025 08:55:24 +0530 Subject: [PATCH 101/123] support latest few-step wan LoRA. (#12541) * support latest few-step wan LoRA. * up * up --- .../loaders/lora_conversion_utils.py | 30 +++++++++++++++---- 1 file changed, 25 insertions(+), 5 deletions(-) diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index 89afb6529a..099dbfc1d2 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -1977,14 +1977,34 @@ def _convert_non_diffusers_wan_lora_to_diffusers(state_dict): "time_projection.1.diff_b" ) - if any("head.head" in k for k in state_dict): - converted_state_dict["proj_out.lora_A.weight"] = original_state_dict.pop( - f"head.head.{lora_down_key}.weight" - ) - converted_state_dict["proj_out.lora_B.weight"] = original_state_dict.pop(f"head.head.{lora_up_key}.weight") + if any("head.head" in k for k in original_state_dict): + if any(f"head.head.{lora_down_key}.weight" in k for k in state_dict): + converted_state_dict["proj_out.lora_A.weight"] = original_state_dict.pop( + f"head.head.{lora_down_key}.weight" + ) + if any(f"head.head.{lora_up_key}.weight" in k for k in state_dict): + converted_state_dict["proj_out.lora_B.weight"] = original_state_dict.pop( + f"head.head.{lora_up_key}.weight" + ) if "head.head.diff_b" in original_state_dict: converted_state_dict["proj_out.lora_B.bias"] = original_state_dict.pop("head.head.diff_b") + # Notes: https://huggingface.co/lightx2v/Wan2.2-Distill-Loras + # This is my (sayakpaul) assumption that this particular key belongs to the down matrix. + # Since for this particular LoRA, we don't have the corresponding up matrix, I will use + # an identity. + if any("head.head" in k and k.endswith(".diff") for k in state_dict): + if f"head.head.{lora_down_key}.weight" in state_dict: + logger.info( + f"The state dict seems to be have both `head.head.diff` and `head.head.{lora_down_key}.weight` keys, which is unexpected." + ) + converted_state_dict["proj_out.lora_A.weight"] = original_state_dict.pop("head.head.diff") + down_matrix_head = converted_state_dict["proj_out.lora_A.weight"] + up_matrix_shape = (down_matrix_head.shape[0], converted_state_dict["proj_out.lora_B.bias"].shape[0]) + converted_state_dict["proj_out.lora_B.weight"] = torch.eye( + *up_matrix_shape, dtype=down_matrix_head.dtype, device=down_matrix_head.device + ).T + for text_time in ["text_embedding", "time_embedding"]: if any(text_time in k for k in original_state_dict): for b_n in [0, 2]: From ecfbc8f952fb20b47d2f85695e98d64c2f5b219a Mon Sep 17 00:00:00 2001 From: Dhruv Nair Date: Tue, 28 Oct 2025 09:21:31 +0530 Subject: [PATCH 102/123] [Pipelines] Enable Wan VACE to run since single transformer (#12428) * update * update * update * update * update --- .../pipelines/wan/pipeline_wan_vace.py | 73 ++++++++++----- tests/pipelines/wan/test_wan_vace.py | 89 ++++++++++++++++++- 2 files changed, 137 insertions(+), 25 deletions(-) diff --git a/src/diffusers/pipelines/wan/pipeline_wan_vace.py b/src/diffusers/pipelines/wan/pipeline_wan_vace.py index 2b1890afec..63e557a98f 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan_vace.py +++ b/src/diffusers/pipelines/wan/pipeline_wan_vace.py @@ -152,34 +152,36 @@ class WanVACEPipeline(DiffusionPipeline, WanLoraLoaderMixin): text_encoder ([`T5EncoderModel`]): [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically the [google/umt5-xxl](https://huggingface.co/google/umt5-xxl) variant. - transformer ([`WanVACETransformer3DModel`]): - Conditional Transformer to denoise the input latents. - transformer_2 ([`WanVACETransformer3DModel`], *optional*): - Conditional Transformer to denoise the input latents during the low-noise stage. In two-stage denoising, - `transformer` handles high-noise stages and `transformer_2` handles low-noise stages. If not provided, only - `transformer` is used. - scheduler ([`UniPCMultistepScheduler`]): - A scheduler to be used in combination with `transformer` to denoise the encoded image latents. vae ([`AutoencoderKLWan`]): Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations. + scheduler ([`UniPCMultistepScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded image latents. + transformer ([`WanVACETransformer3DModel`], *optional*): + Conditional Transformer to denoise the input latents during the high-noise stage. In two-stage denoising, + `transformer` handles high-noise stages and `transformer_2` handles low-noise stages. At least one of + `transformer` or `transformer_2` must be provided. + transformer_2 ([`WanVACETransformer3DModel`], *optional*): + Conditional Transformer to denoise the input latents during the low-noise stage. In two-stage denoising, + `transformer` handles high-noise stages and `transformer_2` handles low-noise stages. At least one of + `transformer` or `transformer_2` must be provided. boundary_ratio (`float`, *optional*, defaults to `None`): Ratio of total timesteps to use as the boundary for switching between transformers in two-stage denoising. The actual boundary timestep is calculated as `boundary_ratio * num_train_timesteps`. When provided, `transformer` handles timesteps >= boundary_timestep and `transformer_2` handles timesteps < - boundary_timestep. If `None`, only `transformer` is used for the entire denoising process. + boundary_timestep. If `None`, only the available transformer is used for the entire denoising process. """ - model_cpu_offload_seq = "text_encoder->transformer->vae" + model_cpu_offload_seq = "text_encoder->transformer->transformer_2->vae" _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - _optional_components = ["transformer_2"] + _optional_components = ["transformer", "transformer_2"] def __init__( self, tokenizer: AutoTokenizer, text_encoder: UMT5EncoderModel, - transformer: WanVACETransformer3DModel, vae: AutoencoderKLWan, scheduler: FlowMatchEulerDiscreteScheduler, + transformer: WanVACETransformer3DModel = None, transformer_2: WanVACETransformer3DModel = None, boundary_ratio: Optional[float] = None, ): @@ -336,7 +338,15 @@ class WanVACEPipeline(DiffusionPipeline, WanLoraLoaderMixin): reference_images=None, guidance_scale_2=None, ): - base = self.vae_scale_factor_spatial * self.transformer.config.patch_size[1] + if self.transformer is not None: + base = self.vae_scale_factor_spatial * self.transformer.config.patch_size[1] + elif self.transformer_2 is not None: + base = self.vae_scale_factor_spatial * self.transformer_2.config.patch_size[1] + else: + raise ValueError( + "`transformer` or `transformer_2` component must be set in order to run inference with this pipeline" + ) + if height % base != 0 or width % base != 0: raise ValueError(f"`height` and `width` have to be divisible by {base} but are {height} and {width}.") @@ -414,7 +424,11 @@ class WanVACEPipeline(DiffusionPipeline, WanLoraLoaderMixin): device: Optional[torch.device] = None, ): if video is not None: - base = self.vae_scale_factor_spatial * self.transformer.config.patch_size[1] + base = self.vae_scale_factor_spatial * ( + self.transformer.config.patch_size[1] + if self.transformer is not None + else self.transformer_2.config.patch_size[1] + ) video_height, video_width = self.video_processor.get_default_height_width(video[0]) if video_height * video_width > height * width: @@ -589,7 +603,11 @@ class WanVACEPipeline(DiffusionPipeline, WanLoraLoaderMixin): "Generating with more than one video is not yet supported. This may be supported in the future." ) - transformer_patch_size = self.transformer.config.patch_size[1] + transformer_patch_size = ( + self.transformer.config.patch_size[1] + if self.transformer is not None + else self.transformer_2.config.patch_size[1] + ) mask_list = [] for mask_, reference_images_batch in zip(mask, reference_images): @@ -844,20 +862,25 @@ class WanVACEPipeline(DiffusionPipeline, WanLoraLoaderMixin): batch_size = prompt_embeds.shape[0] vae_dtype = self.vae.dtype - transformer_dtype = self.transformer.dtype + transformer_dtype = self.transformer.dtype if self.transformer is not None else self.transformer_2.dtype + vace_layers = ( + self.transformer.config.vace_layers + if self.transformer is not None + else self.transformer_2.config.vace_layers + ) if isinstance(conditioning_scale, (int, float)): - conditioning_scale = [conditioning_scale] * len(self.transformer.config.vace_layers) + conditioning_scale = [conditioning_scale] * len(vace_layers) if isinstance(conditioning_scale, list): - if len(conditioning_scale) != len(self.transformer.config.vace_layers): + if len(conditioning_scale) != len(vace_layers): raise ValueError( - f"Length of `conditioning_scale` {len(conditioning_scale)} does not match number of layers {len(self.transformer.config.vace_layers)}." + f"Length of `conditioning_scale` {len(conditioning_scale)} does not match number of layers {len(vace_layers)}." ) conditioning_scale = torch.tensor(conditioning_scale) if isinstance(conditioning_scale, torch.Tensor): - if conditioning_scale.size(0) != len(self.transformer.config.vace_layers): + if conditioning_scale.size(0) != len(vace_layers): raise ValueError( - f"Length of `conditioning_scale` {conditioning_scale.size(0)} does not match number of layers {len(self.transformer.config.vace_layers)}." + f"Length of `conditioning_scale` {conditioning_scale.size(0)} does not match number of layers {len(vace_layers)}." ) conditioning_scale = conditioning_scale.to(device=device, dtype=transformer_dtype) @@ -900,7 +923,11 @@ class WanVACEPipeline(DiffusionPipeline, WanLoraLoaderMixin): conditioning_latents = torch.cat([conditioning_latents, mask], dim=1) conditioning_latents = conditioning_latents.to(transformer_dtype) - num_channels_latents = self.transformer.config.in_channels + num_channels_latents = ( + self.transformer.config.in_channels + if self.transformer is not None + else self.transformer_2.config.in_channels + ) latents = self.prepare_latents( batch_size * num_videos_per_prompt, num_channels_latents, @@ -968,7 +995,7 @@ class WanVACEPipeline(DiffusionPipeline, WanLoraLoaderMixin): attention_kwargs=attention_kwargs, return_dict=False, )[0] - noise_pred = noise_uncond + guidance_scale * (noise_pred - noise_uncond) + noise_pred = noise_uncond + current_guidance_scale * (noise_pred - noise_uncond) # compute the previous noisy sample x_t -> x_t-1 latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] diff --git a/tests/pipelines/wan/test_wan_vace.py b/tests/pipelines/wan/test_wan_vace.py index f99863c880..fe078c0deb 100644 --- a/tests/pipelines/wan/test_wan_vace.py +++ b/tests/pipelines/wan/test_wan_vace.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import tempfile import unittest import numpy as np @@ -19,9 +20,15 @@ import torch from PIL import Image from transformers import AutoTokenizer, T5EncoderModel -from diffusers import AutoencoderKLWan, FlowMatchEulerDiscreteScheduler, WanVACEPipeline, WanVACETransformer3DModel +from diffusers import ( + AutoencoderKLWan, + FlowMatchEulerDiscreteScheduler, + UniPCMultistepScheduler, + WanVACEPipeline, + WanVACETransformer3DModel, +) -from ...testing_utils import enable_full_determinism +from ...testing_utils import enable_full_determinism, torch_device from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS from ..test_pipelines_common import PipelineTesterMixin @@ -212,3 +219,81 @@ class WanVACEPipelineFastTests(PipelineTesterMixin, unittest.TestCase): ) def test_save_load_float16(self): pass + + def test_inference_with_only_transformer(self): + components = self.get_dummy_components() + components["transformer_2"] = None + components["boundary_ratio"] = 0.0 + pipe = self.pipeline_class(**components) + pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(torch_device) + video = pipe(**inputs).frames[0] + assert video.shape == (17, 3, 16, 16) + + def test_inference_with_only_transformer_2(self): + components = self.get_dummy_components() + components["transformer_2"] = components["transformer"] + components["transformer"] = None + + # FlowMatchEulerDiscreteScheduler doesn't support running low noise only scheduler + # because starting timestep t == 1000 == boundary_timestep + components["scheduler"] = UniPCMultistepScheduler( + prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=3.0 + ) + + components["boundary_ratio"] = 1.0 + pipe = self.pipeline_class(**components) + pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(torch_device) + video = pipe(**inputs).frames[0] + assert video.shape == (17, 3, 16, 16) + + def test_save_load_optional_components(self, expected_max_difference=1e-4): + optional_component = ["transformer"] + + components = self.get_dummy_components() + components["transformer_2"] = components["transformer"] + # FlowMatchEulerDiscreteScheduler doesn't support running low noise only scheduler + # because starting timestep t == 1000 == boundary_timestep + components["scheduler"] = UniPCMultistepScheduler( + prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=3.0 + ) + for component in optional_component: + components[component] = None + + components["boundary_ratio"] = 1.0 + + pipe = self.pipeline_class(**components) + for component in pipe.components.values(): + if hasattr(component, "set_default_attn_processor"): + component.set_default_attn_processor() + pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + + generator_device = "cpu" + inputs = self.get_dummy_inputs(generator_device) + torch.manual_seed(0) + output = pipe(**inputs)[0] + + with tempfile.TemporaryDirectory() as tmpdir: + pipe.save_pretrained(tmpdir, safe_serialization=False) + pipe_loaded = self.pipeline_class.from_pretrained(tmpdir) + for component in pipe_loaded.components.values(): + if hasattr(component, "set_default_attn_processor"): + component.set_default_attn_processor() + pipe_loaded.to(torch_device) + pipe_loaded.set_progress_bar_config(disable=None) + + for component in optional_component: + assert getattr(pipe_loaded, component) is None, f"`{component}` did not stay set to None after loading." + + inputs = self.get_dummy_inputs(generator_device) + torch.manual_seed(0) + output_loaded = pipe_loaded(**inputs)[0] + + max_diff = np.abs(output.detach().cpu().numpy() - output_loaded.detach().cpu().numpy()).max() + assert max_diff < expected_max_difference, "Outputs exceed expecpted maximum difference" From dc622a95d0c8dcd1a0badbe6cf6e61ef4eca3917 Mon Sep 17 00:00:00 2001 From: "Wang, Yi" Date: Tue, 28 Oct 2025 11:59:20 +0800 Subject: [PATCH 103/123] fix crash if tiling mode is enabled (#12521) * fix crash in tiling mode is enabled Signed-off-by: Wang, Yi A * fmt Signed-off-by: Wang, Yi A --------- Signed-off-by: Wang, Yi A Co-authored-by: Sayak Paul --- .../models/autoencoders/autoencoder_kl_wan.py | 27 ++++++++++++++----- 1 file changed, 21 insertions(+), 6 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_wan.py b/src/diffusers/models/autoencoders/autoencoder_kl_wan.py index dae26f8086..f8bdfeb755 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_wan.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_wan.py @@ -1337,9 +1337,18 @@ class AutoencoderKLWan(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalMo tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio - - blend_height = self.tile_sample_min_height - self.tile_sample_stride_height - blend_width = self.tile_sample_min_width - self.tile_sample_stride_width + tile_sample_stride_height = self.tile_sample_stride_height + tile_sample_stride_width = self.tile_sample_stride_width + if self.config.patch_size is not None: + sample_height = sample_height // self.config.patch_size + sample_width = sample_width // self.config.patch_size + tile_sample_stride_height = tile_sample_stride_height // self.config.patch_size + tile_sample_stride_width = tile_sample_stride_width // self.config.patch_size + blend_height = self.tile_sample_min_height // self.config.patch_size - tile_sample_stride_height + blend_width = self.tile_sample_min_width // self.config.patch_size - tile_sample_stride_width + else: + blend_height = self.tile_sample_min_height - tile_sample_stride_height + blend_width = self.tile_sample_min_width - tile_sample_stride_width # Split z into overlapping tiles and decode them separately. # The tiles have an overlap to avoid seams between tiles. @@ -1353,7 +1362,9 @@ class AutoencoderKLWan(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalMo self._conv_idx = [0] tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width] tile = self.post_quant_conv(tile) - decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx) + decoded = self.decoder( + tile, feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=(k == 0) + ) time.append(decoded) row.append(torch.cat(time, dim=2)) rows.append(row) @@ -1369,11 +1380,15 @@ class AutoencoderKLWan(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalMo tile = self.blend_v(rows[i - 1][j], tile, blend_height) if j > 0: tile = self.blend_h(row[j - 1], tile, blend_width) - result_row.append(tile[:, :, :, : self.tile_sample_stride_height, : self.tile_sample_stride_width]) + result_row.append(tile[:, :, :, :tile_sample_stride_height, :tile_sample_stride_width]) result_rows.append(torch.cat(result_row, dim=-1)) - dec = torch.cat(result_rows, dim=3)[:, :, :, :sample_height, :sample_width] + if self.config.patch_size is not None: + dec = unpatchify(dec, patch_size=self.config.patch_size) + + dec = torch.clamp(dec, min=-1.0, max=1.0) + if not return_dict: return (dec,) return DecoderOutput(sample=dec) From 40528e9ae7d56740c00d838299198d34111717bb Mon Sep 17 00:00:00 2001 From: Meatfucker <74834323+Meatfucker@users.noreply.github.com> Date: Tue, 28 Oct 2025 01:54:24 -0400 Subject: [PATCH 104/123] Fix typos in kandinsky5 docs (#12552) Update kandinsky5.md Fix typos --- docs/source/en/api/pipelines/kandinsky5.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/source/en/api/pipelines/kandinsky5.md b/docs/source/en/api/pipelines/kandinsky5.md index cb1c119f80..a98a0826b7 100644 --- a/docs/source/en/api/pipelines/kandinsky5.md +++ b/docs/source/en/api/pipelines/kandinsky5.md @@ -92,7 +92,7 @@ pipe = pipe.to("cuda") pipe.transformer.set_attention_backend( "flex" -) # <--- Sett attention bakend to Flex +) # <--- Set attention backend to Flex pipe.transformer.compile( mode="max-autotune-no-cudagraphs", dynamic=True @@ -115,7 +115,7 @@ export_to_video(output, "output.mp4", fps=24, quality=9) ``` ### Diffusion Distilled model -**⚠️ Warning!** all nocfg and diffusion distilled models should be infered wothout CFG (```guidance_scale=1.0```): +**⚠️ Warning!** all nocfg and diffusion distilled models should be inferred without CFG (```guidance_scale=1.0```): ```python model_id = "ai-forever/Kandinsky-5.0-T2V-Lite-distilled16steps-5s-Diffusers" @@ -146,4 +146,4 @@ export_to_video(output, "output.mp4", fps=24, quality=9) howpublished = {\url{https://github.com/ai-forever/Kandinsky-5}}, year = 2025 } -``` \ No newline at end of file +``` From 55d49d4379007740af20629bb61aba9546c6b053 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Tue, 28 Oct 2025 13:29:51 +0530 Subject: [PATCH 105/123] [ci] don't run sana layerwise casting tests in CI. (#12551) * don't run sana layerwise casting tests in CI. * up --- tests/lora/test_lora_layers_sana.py | 6 +++++- .../models/autoencoders/test_models_autoencoder_dc.py | 10 +++++----- tests/pipelines/sana/test_sana.py | 5 +++++ tests/pipelines/sana/test_sana_controlnet.py | 9 +++++---- tests/pipelines/sana/test_sana_sprint.py | 9 +++++---- tests/pipelines/sana/test_sana_sprint_img2img.py | 9 +++++---- tests/testing_utils.py | 2 ++ 7 files changed, 32 insertions(+), 18 deletions(-) diff --git a/tests/lora/test_lora_layers_sana.py b/tests/lora/test_lora_layers_sana.py index 3cdb28de75..a860b7b44f 100644 --- a/tests/lora/test_lora_layers_sana.py +++ b/tests/lora/test_lora_layers_sana.py @@ -20,7 +20,7 @@ from transformers import Gemma2Model, GemmaTokenizer from diffusers import AutoencoderDC, FlowMatchEulerDiscreteScheduler, SanaPipeline, SanaTransformer2DModel -from ..testing_utils import floats_tensor, require_peft_backend +from ..testing_utils import IS_GITHUB_ACTIONS, floats_tensor, require_peft_backend sys.path.append(".") @@ -136,3 +136,7 @@ class SanaLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): @unittest.skip("Text encoder LoRA is not supported in SANA.") def test_simple_inference_with_text_lora_save_load(self): pass + + @unittest.skipIf(IS_GITHUB_ACTIONS, reason="Skipping test inside GitHub Actions environment") + def test_layerwise_casting_inference_denoiser(self): + return super().test_layerwise_casting_inference_denoiser() diff --git a/tests/models/autoencoders/test_models_autoencoder_dc.py b/tests/models/autoencoders/test_models_autoencoder_dc.py index a6912f3eba..d34001e7b9 100644 --- a/tests/models/autoencoders/test_models_autoencoder_dc.py +++ b/tests/models/autoencoders/test_models_autoencoder_dc.py @@ -17,11 +17,7 @@ import unittest from diffusers import AutoencoderDC -from ...testing_utils import ( - enable_full_determinism, - floats_tensor, - torch_device, -) +from ...testing_utils import IS_GITHUB_ACTIONS, enable_full_determinism, floats_tensor, torch_device from ..test_modeling_common import ModelTesterMixin from .testing_utils import AutoencoderTesterMixin @@ -82,3 +78,7 @@ class AutoencoderDCTests(ModelTesterMixin, AutoencoderTesterMixin, unittest.Test init_dict = self.get_autoencoder_dc_config() inputs_dict = self.dummy_input return init_dict, inputs_dict + + @unittest.skipIf(IS_GITHUB_ACTIONS, reason="Skipping test inside GitHub Actions environment") + def test_layerwise_casting_inference(self): + super().test_layerwise_casting_inference() diff --git a/tests/pipelines/sana/test_sana.py b/tests/pipelines/sana/test_sana.py index 34ea3079b1..f23303c966 100644 --- a/tests/pipelines/sana/test_sana.py +++ b/tests/pipelines/sana/test_sana.py @@ -23,6 +23,7 @@ from transformers import Gemma2Config, Gemma2Model, GemmaTokenizer from diffusers import AutoencoderDC, FlowMatchEulerDiscreteScheduler, SanaPipeline, SanaTransformer2DModel from ...testing_utils import ( + IS_GITHUB_ACTIONS, backend_empty_cache, enable_full_determinism, require_torch_accelerator, @@ -304,6 +305,10 @@ class SanaPipelineFastTests(PipelineTesterMixin, unittest.TestCase): # Requires higher tolerance as model seems very sensitive to dtype super().test_float16_inference(expected_max_diff=0.08) + @unittest.skipIf(IS_GITHUB_ACTIONS, reason="Skipping test inside GitHub Actions environment") + def test_layerwise_casting_inference(self): + super().test_layerwise_casting_inference() + @slow @require_torch_accelerator diff --git a/tests/pipelines/sana/test_sana_controlnet.py b/tests/pipelines/sana/test_sana_controlnet.py index 043e276fcb..df14d935ed 100644 --- a/tests/pipelines/sana/test_sana_controlnet.py +++ b/tests/pipelines/sana/test_sana_controlnet.py @@ -28,10 +28,7 @@ from diffusers import ( ) from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import ( - enable_full_determinism, - torch_device, -) +from ...testing_utils import IS_GITHUB_ACTIONS, enable_full_determinism, torch_device from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS from ..test_pipelines_common import PipelineTesterMixin, to_np @@ -326,3 +323,7 @@ class SanaControlNetPipelineFastTests(PipelineTesterMixin, unittest.TestCase): def test_float16_inference(self): # Requires higher tolerance as model seems very sensitive to dtype super().test_float16_inference(expected_max_diff=0.08) + + @unittest.skipIf(IS_GITHUB_ACTIONS, reason="Skipping test inside GitHub Actions environment") + def test_layerwise_casting_inference(self): + super().test_layerwise_casting_inference() diff --git a/tests/pipelines/sana/test_sana_sprint.py b/tests/pipelines/sana/test_sana_sprint.py index fee2304dce..0d45205ea8 100644 --- a/tests/pipelines/sana/test_sana_sprint.py +++ b/tests/pipelines/sana/test_sana_sprint.py @@ -21,10 +21,7 @@ from transformers import Gemma2Config, Gemma2Model, GemmaTokenizer from diffusers import AutoencoderDC, SanaSprintPipeline, SanaTransformer2DModel, SCMScheduler -from ...testing_utils import ( - enable_full_determinism, - torch_device, -) +from ...testing_utils import IS_GITHUB_ACTIONS, enable_full_determinism, torch_device from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS from ..test_pipelines_common import PipelineTesterMixin, to_np @@ -300,3 +297,7 @@ class SanaSprintPipelineFastTests(PipelineTesterMixin, unittest.TestCase): def test_float16_inference(self): # Requires higher tolerance as model seems very sensitive to dtype super().test_float16_inference(expected_max_diff=0.08) + + @unittest.skipIf(IS_GITHUB_ACTIONS, reason="Skipping test inside GitHub Actions environment") + def test_layerwise_casting_inference(self): + super().test_layerwise_casting_inference() diff --git a/tests/pipelines/sana/test_sana_sprint_img2img.py b/tests/pipelines/sana/test_sana_sprint_img2img.py index c218abb8e9..5de5c7f446 100644 --- a/tests/pipelines/sana/test_sana_sprint_img2img.py +++ b/tests/pipelines/sana/test_sana_sprint_img2img.py @@ -22,10 +22,7 @@ from transformers import Gemma2Config, Gemma2Model, GemmaTokenizer from diffusers import AutoencoderDC, SanaSprintImg2ImgPipeline, SanaTransformer2DModel, SCMScheduler from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import ( - enable_full_determinism, - torch_device, -) +from ...testing_utils import IS_GITHUB_ACTIONS, enable_full_determinism, torch_device from ..pipeline_params import ( IMAGE_TO_IMAGE_IMAGE_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, @@ -312,3 +309,7 @@ class SanaSprintImg2ImgPipelineFastTests(PipelineTesterMixin, unittest.TestCase) def test_float16_inference(self): # Requires higher tolerance as model seems very sensitive to dtype super().test_float16_inference(expected_max_diff=0.08) + + @unittest.skipIf(IS_GITHUB_ACTIONS, reason="Skipping test inside GitHub Actions environment") + def test_layerwise_casting_inference(self): + super().test_layerwise_casting_inference() diff --git a/tests/testing_utils.py b/tests/testing_utils.py index 7f849219c1..951ba41280 100644 --- a/tests/testing_utils.py +++ b/tests/testing_utils.py @@ -63,6 +63,8 @@ else: IS_CUDA_SYSTEM = False IS_XPU_SYSTEM = False +IS_GITHUB_ACTIONS = os.getenv("GITHUB_ACTIONS") == "true" and os.getenv("DIFFUSERS_IS_CI") == "yes" + global_rng = random.Random() logger = get_logger(__name__) From 84e16575e4c5e90b6b49301cfa162ced4cf478d2 Mon Sep 17 00:00:00 2001 From: galbria <158810732+galbria@users.noreply.github.com> Date: Tue, 28 Oct 2025 12:57:48 +0200 Subject: [PATCH 106/123] Bria fibo (#12545) * Bria FIBO pipeline * style fixs * fix CR * Refactor BriaFibo classes and update pipeline parameters - Updated BriaFiboAttnProcessor and BriaFiboAttention classes to reflect changes from Flux equivalents. - Modified the _unpack_latents method in BriaFiboPipeline to improve clarity. - Increased the default max_sequence_length to 3000 and added a new optional parameter do_patching. - Cleaned up test_pipeline_bria_fibo.py by removing unused imports and skipping unsupported tests. * edit the docs of FIBO * Remove unused BriaFibo imports and update CPU offload method in BriaFiboPipeline * Refactor FIBO classes to BriaFibo naming convention - Updated class names from FIBO to BriaFibo for consistency across the module. - Modified instances of FIBOEmbedND, FIBOTimesteps, TextProjection, and TimestepProjEmbeddings to reflect the new naming. - Ensured all references in the BriaFiboTransformer2DModel are updated accordingly. * Add BriaFiboTransformer2DModel import to transformers module * Remove unused BriaFibo imports from modular pipelines and add BriaFiboTransformer2DModel and BriaFiboPipeline classes to dummy objects for enhanced compatibility with torch and transformers. * Update BriaFibo classes with copied documentation and fix import typo in pipeline module - Added documentation comments indicating the source of copied code in BriaFiboTransformerBlock and _pack_latents methods. - Corrected the import statement for BriaFiboPipeline in the pipelines module. * Remove unused BriaFibo imports from __init__.py to streamline modular pipelines. * Refactor documentation comments in BriaFibo classes to indicate inspiration from existing implementations - Updated comments in BriaFiboAttnProcessor, BriaFiboAttention, and BriaFiboPipeline to reflect that the code is inspired by other modules rather than copied. - Enhanced clarity on the origins of the methods to maintain proper attribution. * change Inspired by to Based on * add reference link and fix trailing whitespace * Add BriaFiboTransformer2DModel documentation and update comments in BriaFibo classes - Introduced a new documentation file for BriaFiboTransformer2DModel. - Updated comments in BriaFiboAttnProcessor, BriaFiboAttention, and BriaFiboPipeline to clarify the origins of the code, indicating copied sources for better attribution. --------- Co-authored-by: sayakpaul --- docs/source/en/_toctree.yml | 4 + .../en/api/models/transformer_bria_fibo.md | 19 + docs/source/en/api/pipelines/bria_fibo.md | 45 + src/diffusers/__init__.py | 4 + src/diffusers/models/__init__.py | 2 + src/diffusers/models/transformers/__init__.py | 1 + .../transformers/transformer_bria_fibo.py | 655 ++++++++++++++ src/diffusers/pipelines/__init__.py | 2 + src/diffusers/pipelines/bria_fibo/__init__.py | 48 + .../pipelines/bria_fibo/pipeline_bria_fibo.py | 838 ++++++++++++++++++ .../pipelines/bria_fibo/pipeline_output.py | 21 + src/diffusers/utils/dummy_pt_objects.py | 15 + .../dummy_torch_and_transformers_objects.py | 15 + .../test_models_transformer_bria_fibo.py | 89 ++ tests/pipelines/bria_fibo/__init__.py | 0 .../bria_fibo/test_pipeline_bria_fibo.py | 139 +++ 16 files changed, 1897 insertions(+) create mode 100644 docs/source/en/api/models/transformer_bria_fibo.md create mode 100644 docs/source/en/api/pipelines/bria_fibo.md create mode 100644 src/diffusers/models/transformers/transformer_bria_fibo.py create mode 100644 src/diffusers/pipelines/bria_fibo/__init__.py create mode 100644 src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo.py create mode 100644 src/diffusers/pipelines/bria_fibo/pipeline_output.py create mode 100644 tests/models/transformers/test_models_transformer_bria_fibo.py create mode 100644 tests/pipelines/bria_fibo/__init__.py create mode 100644 tests/pipelines/bria_fibo/test_pipeline_bria_fibo.py diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 44870f680e..8103c01643 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -323,6 +323,8 @@ title: AllegroTransformer3DModel - local: api/models/aura_flow_transformer2d title: AuraFlowTransformer2DModel + - local: api/models/transformer_bria_fibo + title: BriaFiboTransformer2DModel - local: api/models/bria_transformer title: BriaTransformer2DModel - local: api/models/chroma_transformer @@ -469,6 +471,8 @@ title: BLIP-Diffusion - local: api/pipelines/bria_3_2 title: Bria 3.2 + - local: api/pipelines/bria_fibo + title: Bria Fibo - local: api/pipelines/chroma title: Chroma - local: api/pipelines/cogview3 diff --git a/docs/source/en/api/models/transformer_bria_fibo.md b/docs/source/en/api/models/transformer_bria_fibo.md new file mode 100644 index 0000000000..5691746ccd --- /dev/null +++ b/docs/source/en/api/models/transformer_bria_fibo.md @@ -0,0 +1,19 @@ + + +# BriaFiboTransformer2DModel + +A modified flux Transformer model from [Bria](https://huggingface.co/briaai/FIBO) + +## BriaFiboTransformer2DModel + +[[autodoc]] BriaFiboTransformer2DModel diff --git a/docs/source/en/api/pipelines/bria_fibo.md b/docs/source/en/api/pipelines/bria_fibo.md new file mode 100644 index 0000000000..96cad10dda --- /dev/null +++ b/docs/source/en/api/pipelines/bria_fibo.md @@ -0,0 +1,45 @@ + + +# Bria Fibo + +Text-to-image models have mastered imagination - but not control. FIBO changes that. + +FIBO is trained on structured JSON captions up to 1,000+ words and designed to understand and control different visual parameters such as lighting, composition, color, and camera settings, enabling precise and reproducible outputs. + +With only 8 billion parameters, FIBO provides a new level of image quality, prompt adherence and proffesional control. + +FIBO is trained exclusively on a structured prompt and will not work with freeform text prompts. +you can use the [FIBO-VLM-prompt-to-JSON](https://huggingface.co/briaai/FIBO-VLM-prompt-to-JSON) model or the [FIBO-gemini-prompt-to-JSON](https://huggingface.co/briaai/FIBO-gemini-prompt-to-JSON) to convert your freeform text prompt to a structured JSON prompt. + +its not recommended to use freeform text prompts directly with FIBO, as it will not produce the best results. + +you can learn more about FIBO in [Bria Fibo Hugging Face page](https://huggingface.co/briaai/FIBO). + + +## Usage + +_As the model is gated, before using it with diffusers you first need to go to the [Bria Fibo Hugging Face page](https://huggingface.co/briaai/FIBO), fill in the form and accept the gate. Once you are in, you need to login so that your system knows you’ve accepted the gate._ + +Use the command below to log in: + +```bash +hf auth login +``` + + +## BriaPipeline + +[[autodoc]] BriaPipeline + - all + - __call__ + diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 2a8b589846..94104667b5 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -198,6 +198,7 @@ else: "AutoencoderOobleck", "AutoencoderTiny", "AutoModel", + "BriaFiboTransformer2DModel", "BriaTransformer2DModel", "CacheMixin", "ChromaTransformer2DModel", @@ -430,6 +431,7 @@ else: "AuraFlowPipeline", "BlipDiffusionControlNetPipeline", "BlipDiffusionPipeline", + "BriaFiboPipeline", "BriaPipeline", "ChromaImg2ImgPipeline", "ChromaPipeline", @@ -901,6 +903,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: AutoencoderOobleck, AutoencoderTiny, AutoModel, + BriaFiboTransformer2DModel, BriaTransformer2DModel, CacheMixin, ChromaTransformer2DModel, @@ -1103,6 +1106,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: AudioLDM2UNet2DConditionModel, AudioLDMPipeline, AuraFlowPipeline, + BriaFiboPipeline, BriaPipeline, ChromaImg2ImgPipeline, ChromaPipeline, diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index 241ccf7b78..e3b2974641 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -84,6 +84,7 @@ if is_torch_available(): _import_structure["transformers.transformer_2d"] = ["Transformer2DModel"] _import_structure["transformers.transformer_allegro"] = ["AllegroTransformer3DModel"] _import_structure["transformers.transformer_bria"] = ["BriaTransformer2DModel"] + _import_structure["transformers.transformer_bria_fibo"] = ["BriaFiboTransformer2DModel"] _import_structure["transformers.transformer_chroma"] = ["ChromaTransformer2DModel"] _import_structure["transformers.transformer_cogview3plus"] = ["CogView3PlusTransformer2DModel"] _import_structure["transformers.transformer_cogview4"] = ["CogView4Transformer2DModel"] @@ -174,6 +175,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: from .transformers import ( AllegroTransformer3DModel, AuraFlowTransformer2DModel, + BriaFiboTransformer2DModel, BriaTransformer2DModel, ChromaTransformer2DModel, CogVideoXTransformer3DModel, diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index 6c66131dea..2fe1159eec 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -18,6 +18,7 @@ if is_torch_available(): from .transformer_2d import Transformer2DModel from .transformer_allegro import AllegroTransformer3DModel from .transformer_bria import BriaTransformer2DModel + from .transformer_bria_fibo import BriaFiboTransformer2DModel from .transformer_chroma import ChromaTransformer2DModel from .transformer_cogview3plus import CogView3PlusTransformer2DModel from .transformer_cogview4 import CogView4Transformer2DModel diff --git a/src/diffusers/models/transformers/transformer_bria_fibo.py b/src/diffusers/models/transformers/transformer_bria_fibo.py new file mode 100644 index 0000000000..09f7961932 --- /dev/null +++ b/src/diffusers/models/transformers/transformer_bria_fibo.py @@ -0,0 +1,655 @@ +# Copyright (c) Bria.ai. All rights reserved. +# +# This file is licensed under the Creative Commons Attribution-NonCommercial 4.0 International Public License (CC-BY-NC-4.0). +# You may obtain a copy of the license at https://creativecommons.org/licenses/by-nc/4.0/ +# +# You are free to share and adapt this material for non-commercial purposes provided you give appropriate credit, +# indicate if changes were made, and do not use the material for commercial purposes. +# +# See the license for further details. +import inspect +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin, PeftAdapterMixin +from ...models.attention_processor import Attention +from ...models.embeddings import TimestepEmbedding, apply_rotary_emb, get_1d_rotary_pos_embed, get_timestep_embedding +from ...models.modeling_outputs import Transformer2DModelOutput +from ...models.modeling_utils import ModelMixin +from ...models.transformers.transformer_bria import BriaAttnProcessor +from ...utils import ( + USE_PEFT_BACKEND, + logging, + scale_lora_layers, + unscale_lora_layers, +) +from ...utils.torch_utils import maybe_allow_in_graph +from ..attention import AttentionModuleMixin, FeedForward +from ..attention_dispatch import dispatch_attention_fn +from ..normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def _get_projections(attn: "BriaFiboAttention", hidden_states, encoder_hidden_states=None): + query = attn.to_q(hidden_states) + key = attn.to_k(hidden_states) + value = attn.to_v(hidden_states) + + encoder_query = encoder_key = encoder_value = None + if encoder_hidden_states is not None and attn.added_kv_proj_dim is not None: + encoder_query = attn.add_q_proj(encoder_hidden_states) + encoder_key = attn.add_k_proj(encoder_hidden_states) + encoder_value = attn.add_v_proj(encoder_hidden_states) + + return query, key, value, encoder_query, encoder_key, encoder_value + + +def _get_fused_projections(attn: "BriaFiboAttention", hidden_states, encoder_hidden_states=None): + query, key, value = attn.to_qkv(hidden_states).chunk(3, dim=-1) + + encoder_query = encoder_key = encoder_value = (None,) + if encoder_hidden_states is not None and hasattr(attn, "to_added_qkv"): + encoder_query, encoder_key, encoder_value = attn.to_added_qkv(encoder_hidden_states).chunk(3, dim=-1) + + return query, key, value, encoder_query, encoder_key, encoder_value + + +def _get_qkv_projections(attn: "BriaFiboAttention", hidden_states, encoder_hidden_states=None): + if attn.fused_projections: + return _get_fused_projections(attn, hidden_states, encoder_hidden_states) + return _get_projections(attn, hidden_states, encoder_hidden_states) + + +# Copied from diffusers.models.transformers.transformer_flux.FluxAttnProcessor with FluxAttnProcessor->BriaFiboAttnProcessor, FluxAttention->BriaFiboAttention +class BriaFiboAttnProcessor: + _attention_backend = None + _parallel_config = None + + def __init__(self): + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError(f"{self.__class__.__name__} requires PyTorch 2.0. Please upgrade your pytorch version.") + + def __call__( + self, + attn: "BriaFiboAttention", + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor = None, + attention_mask: Optional[torch.Tensor] = None, + image_rotary_emb: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections( + attn, hidden_states, encoder_hidden_states + ) + + query = query.unflatten(-1, (attn.heads, -1)) + key = key.unflatten(-1, (attn.heads, -1)) + value = value.unflatten(-1, (attn.heads, -1)) + + query = attn.norm_q(query) + key = attn.norm_k(key) + + if attn.added_kv_proj_dim is not None: + encoder_query = encoder_query.unflatten(-1, (attn.heads, -1)) + encoder_key = encoder_key.unflatten(-1, (attn.heads, -1)) + encoder_value = encoder_value.unflatten(-1, (attn.heads, -1)) + + encoder_query = attn.norm_added_q(encoder_query) + encoder_key = attn.norm_added_k(encoder_key) + + query = torch.cat([encoder_query, query], dim=1) + key = torch.cat([encoder_key, key], dim=1) + value = torch.cat([encoder_value, value], dim=1) + + if image_rotary_emb is not None: + query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1) + key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1) + + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.to(query.dtype) + + if encoder_hidden_states is not None: + encoder_hidden_states, hidden_states = hidden_states.split_with_sizes( + [encoder_hidden_states.shape[1], hidden_states.shape[1] - encoder_hidden_states.shape[1]], dim=1 + ) + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + encoder_hidden_states = attn.to_add_out(encoder_hidden_states) + + return hidden_states, encoder_hidden_states + else: + return hidden_states + + +# Based on https://github.com/huggingface/diffusers/blob/55d49d4379007740af20629bb61aba9546c6b053/src/diffusers/models/transformers/transformer_flux.py +class BriaFiboAttention(torch.nn.Module, AttentionModuleMixin): + _default_processor_cls = BriaFiboAttnProcessor + _available_processors = [BriaFiboAttnProcessor] + + def __init__( + self, + query_dim: int, + heads: int = 8, + dim_head: int = 64, + dropout: float = 0.0, + bias: bool = False, + added_kv_proj_dim: Optional[int] = None, + added_proj_bias: Optional[bool] = True, + out_bias: bool = True, + eps: float = 1e-5, + out_dim: int = None, + context_pre_only: Optional[bool] = None, + pre_only: bool = False, + elementwise_affine: bool = True, + processor=None, + ): + super().__init__() + + self.head_dim = dim_head + self.inner_dim = out_dim if out_dim is not None else dim_head * heads + self.query_dim = query_dim + self.use_bias = bias + self.dropout = dropout + self.out_dim = out_dim if out_dim is not None else query_dim + self.context_pre_only = context_pre_only + self.pre_only = pre_only + self.heads = out_dim // dim_head if out_dim is not None else heads + self.added_kv_proj_dim = added_kv_proj_dim + self.added_proj_bias = added_proj_bias + + self.norm_q = torch.nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) + self.norm_k = torch.nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) + self.to_q = torch.nn.Linear(query_dim, self.inner_dim, bias=bias) + self.to_k = torch.nn.Linear(query_dim, self.inner_dim, bias=bias) + self.to_v = torch.nn.Linear(query_dim, self.inner_dim, bias=bias) + + if not self.pre_only: + self.to_out = torch.nn.ModuleList([]) + self.to_out.append(torch.nn.Linear(self.inner_dim, self.out_dim, bias=out_bias)) + self.to_out.append(torch.nn.Dropout(dropout)) + + if added_kv_proj_dim is not None: + self.norm_added_q = torch.nn.RMSNorm(dim_head, eps=eps) + self.norm_added_k = torch.nn.RMSNorm(dim_head, eps=eps) + self.add_q_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias) + self.add_k_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias) + self.add_v_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias) + self.to_add_out = torch.nn.Linear(self.inner_dim, query_dim, bias=out_bias) + + if processor is None: + processor = self._default_processor_cls() + self.set_processor(processor) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + image_rotary_emb: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys()) + quiet_attn_parameters = {"ip_adapter_masks", "ip_hidden_states"} + unused_kwargs = [k for k, _ in kwargs.items() if k not in attn_parameters and k not in quiet_attn_parameters] + if len(unused_kwargs) > 0: + logger.warning( + f"joint_attention_kwargs {unused_kwargs} are not expected by {self.processor.__class__.__name__} and will be ignored." + ) + kwargs = {k: w for k, w in kwargs.items() if k in attn_parameters} + return self.processor(self, hidden_states, encoder_hidden_states, attention_mask, image_rotary_emb, **kwargs) + + +class BriaFiboEmbedND(torch.nn.Module): + # modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11 + def __init__(self, theta: int, axes_dim: List[int]): + super().__init__() + self.theta = theta + self.axes_dim = axes_dim + + def forward(self, ids: torch.Tensor) -> torch.Tensor: + n_axes = ids.shape[-1] + cos_out = [] + sin_out = [] + pos = ids.float() + is_mps = ids.device.type == "mps" + freqs_dtype = torch.float32 if is_mps else torch.float64 + for i in range(n_axes): + cos, sin = get_1d_rotary_pos_embed( + self.axes_dim[i], + pos[:, i], + theta=self.theta, + repeat_interleave_real=True, + use_real=True, + freqs_dtype=freqs_dtype, + ) + cos_out.append(cos) + sin_out.append(sin) + freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device) + freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device) + return freqs_cos, freqs_sin + + +@maybe_allow_in_graph +class BriaFiboSingleTransformerBlock(nn.Module): + def __init__(self, dim: int, num_attention_heads: int, attention_head_dim: int, mlp_ratio: float = 4.0): + super().__init__() + self.mlp_hidden_dim = int(dim * mlp_ratio) + + self.norm = AdaLayerNormZeroSingle(dim) + self.proj_mlp = nn.Linear(dim, self.mlp_hidden_dim) + self.act_mlp = nn.GELU(approximate="tanh") + self.proj_out = nn.Linear(dim + self.mlp_hidden_dim, dim) + + processor = BriaAttnProcessor() + + self.attn = Attention( + query_dim=dim, + cross_attention_dim=None, + dim_head=attention_head_dim, + heads=num_attention_heads, + out_dim=dim, + bias=True, + processor=processor, + qk_norm="rms_norm", + eps=1e-6, + pre_only=True, + ) + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + ) -> torch.Tensor: + residual = hidden_states + norm_hidden_states, gate = self.norm(hidden_states, emb=temb) + mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states)) + joint_attention_kwargs = joint_attention_kwargs or {} + attn_output = self.attn( + hidden_states=norm_hidden_states, + image_rotary_emb=image_rotary_emb, + **joint_attention_kwargs, + ) + + hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2) + gate = gate.unsqueeze(1) + hidden_states = gate * self.proj_out(hidden_states) + hidden_states = residual + hidden_states + if hidden_states.dtype == torch.float16: + hidden_states = hidden_states.clip(-65504, 65504) + + return hidden_states + + +class BriaFiboTextProjection(nn.Module): + def __init__(self, in_features, hidden_size): + super().__init__() + self.linear = nn.Linear(in_features=in_features, out_features=hidden_size, bias=False) + + def forward(self, caption): + hidden_states = self.linear(caption) + return hidden_states + + +@maybe_allow_in_graph +# Based on from diffusers.models.transformers.transformer_flux.FluxTransformerBlock +class BriaFiboTransformerBlock(nn.Module): + def __init__( + self, dim: int, num_attention_heads: int, attention_head_dim: int, qk_norm: str = "rms_norm", eps: float = 1e-6 + ): + super().__init__() + + self.norm1 = AdaLayerNormZero(dim) + self.norm1_context = AdaLayerNormZero(dim) + + self.attn = BriaFiboAttention( + query_dim=dim, + added_kv_proj_dim=dim, + dim_head=attention_head_dim, + heads=num_attention_heads, + out_dim=dim, + context_pre_only=False, + bias=True, + processor=BriaFiboAttnProcessor(), + eps=eps, + ) + + self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate") + + self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff_context = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate") + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb) + + norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context( + encoder_hidden_states, emb=temb + ) + joint_attention_kwargs = joint_attention_kwargs or {} + + # Attention. + attention_outputs = self.attn( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_encoder_hidden_states, + image_rotary_emb=image_rotary_emb, + **joint_attention_kwargs, + ) + + if len(attention_outputs) == 2: + attn_output, context_attn_output = attention_outputs + elif len(attention_outputs) == 3: + attn_output, context_attn_output, ip_attn_output = attention_outputs + + # Process attention outputs for the `hidden_states`. + attn_output = gate_msa.unsqueeze(1) * attn_output + hidden_states = hidden_states + attn_output + + norm_hidden_states = self.norm2(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + + ff_output = self.ff(norm_hidden_states) + ff_output = gate_mlp.unsqueeze(1) * ff_output + + hidden_states = hidden_states + ff_output + if len(attention_outputs) == 3: + hidden_states = hidden_states + ip_attn_output + + # Process attention outputs for the `encoder_hidden_states`. + context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output + encoder_hidden_states = encoder_hidden_states + context_attn_output + + norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) + norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] + + context_ff_output = self.ff_context(norm_encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output + if encoder_hidden_states.dtype == torch.float16: + encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504) + + return encoder_hidden_states, hidden_states + + +class BriaFiboTimesteps(nn.Module): + def __init__( + self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1, time_theta=10000 + ): + super().__init__() + self.num_channels = num_channels + self.flip_sin_to_cos = flip_sin_to_cos + self.downscale_freq_shift = downscale_freq_shift + self.scale = scale + self.time_theta = time_theta + + def forward(self, timesteps): + t_emb = get_timestep_embedding( + timesteps, + self.num_channels, + flip_sin_to_cos=self.flip_sin_to_cos, + downscale_freq_shift=self.downscale_freq_shift, + scale=self.scale, + max_period=self.time_theta, + ) + return t_emb + + +class BriaFiboTimestepProjEmbeddings(nn.Module): + def __init__(self, embedding_dim, time_theta): + super().__init__() + + self.time_proj = BriaFiboTimesteps( + num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0, time_theta=time_theta + ) + self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + def forward(self, timestep, dtype): + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=dtype)) # (N, D) + return timesteps_emb + + +class BriaFiboTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin): + """ + Parameters: + patch_size (`int`): Patch size to turn the input data into small patches. + in_channels (`int`, *optional*, defaults to 16): The number of channels in the input. + num_layers (`int`, *optional*, defaults to 18): The number of layers of MMDiT blocks to use. + num_single_layers (`int`, *optional*, defaults to 18): The number of layers of single DiT blocks to use. + attention_head_dim (`int`, *optional*, defaults to 64): The number of channels in each head. + num_attention_heads (`int`, *optional*, defaults to 18): The number of heads to use for multi-head attention. + joint_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use. + pooled_projection_dim (`int`): Number of dimensions to use when projecting the `pooled_projections`. + guidance_embeds (`bool`, defaults to False): Whether to use guidance embeddings. + ... + """ + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + patch_size: int = 1, + in_channels: int = 64, + num_layers: int = 19, + num_single_layers: int = 38, + attention_head_dim: int = 128, + num_attention_heads: int = 24, + joint_attention_dim: int = 4096, + pooled_projection_dim: int = None, + guidance_embeds: bool = False, + axes_dims_rope: List[int] = [16, 56, 56], + rope_theta=10000, + time_theta=10000, + text_encoder_dim: int = 2048, + ): + super().__init__() + self.out_channels = in_channels + self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim + + self.pos_embed = BriaFiboEmbedND(theta=rope_theta, axes_dim=axes_dims_rope) + + self.time_embed = BriaFiboTimestepProjEmbeddings(embedding_dim=self.inner_dim, time_theta=time_theta) + + if guidance_embeds: + self.guidance_embed = BriaFiboTimestepProjEmbeddings(embedding_dim=self.inner_dim) + + self.context_embedder = nn.Linear(self.config.joint_attention_dim, self.inner_dim) + self.x_embedder = torch.nn.Linear(self.config.in_channels, self.inner_dim) + + self.transformer_blocks = nn.ModuleList( + [ + BriaFiboTransformerBlock( + dim=self.inner_dim, + num_attention_heads=self.config.num_attention_heads, + attention_head_dim=self.config.attention_head_dim, + ) + for i in range(self.config.num_layers) + ] + ) + + self.single_transformer_blocks = nn.ModuleList( + [ + BriaFiboSingleTransformerBlock( + dim=self.inner_dim, + num_attention_heads=self.config.num_attention_heads, + attention_head_dim=self.config.attention_head_dim, + ) + for i in range(self.config.num_single_layers) + ] + ) + + self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6) + self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True) + + self.gradient_checkpointing = False + + caption_projection = [ + BriaFiboTextProjection(in_features=text_encoder_dim, hidden_size=self.inner_dim // 2) + for i in range(self.config.num_layers + self.config.num_single_layers) + ] + self.caption_projection = nn.ModuleList(caption_projection) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor = None, + text_encoder_layers: list = None, + pooled_projections: torch.Tensor = None, + timestep: torch.LongTensor = None, + img_ids: torch.Tensor = None, + txt_ids: torch.Tensor = None, + guidance: torch.Tensor = None, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + return_dict: bool = True, + ) -> Union[torch.FloatTensor, Transformer2DModelOutput]: + """ + + Args: + hidden_states (`torch.FloatTensor` of shape `(batch size, channel, height, width)`): + Input `hidden_states`. + encoder_hidden_states (`torch.FloatTensor` of shape `(batch size, sequence_len, embed_dims)`): + Conditional embeddings (embeddings computed from the input conditions such as prompts) to use. + pooled_projections (`torch.FloatTensor` of shape `(batch_size, projection_dim)`): Embeddings projected + from the embeddings of input conditions. + timestep ( `torch.LongTensor`): + Used to indicate denoising step. + joint_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain + tuple. + Returns: + If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a + `tuple` where the first element is the sample tensor. + """ + if joint_attention_kwargs is not None: + joint_attention_kwargs = joint_attention_kwargs.copy() + lora_scale = joint_attention_kwargs.pop("scale", 1.0) + else: + lora_scale = 1.0 + + if USE_PEFT_BACKEND: + # weight the lora layers by setting `lora_scale` for each PEFT layer + scale_lora_layers(self, lora_scale) + else: + if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None: + logger.warning( + "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective." + ) + hidden_states = self.x_embedder(hidden_states) + + timestep = timestep.to(hidden_states.dtype) + if guidance is not None: + guidance = guidance.to(hidden_states.dtype) + else: + guidance = None + + temb = self.time_embed(timestep, dtype=hidden_states.dtype) + + if guidance: + temb += self.guidance_embed(guidance, dtype=hidden_states.dtype) + + encoder_hidden_states = self.context_embedder(encoder_hidden_states) + + if len(txt_ids.shape) == 3: + txt_ids = txt_ids[0] + + if len(img_ids.shape) == 3: + img_ids = img_ids[0] + + ids = torch.cat((txt_ids, img_ids), dim=0) + image_rotary_emb = self.pos_embed(ids) + + new_text_encoder_layers = [] + for i, text_encoder_layer in enumerate(text_encoder_layers): + text_encoder_layer = self.caption_projection[i](text_encoder_layer) + new_text_encoder_layers.append(text_encoder_layer) + text_encoder_layers = new_text_encoder_layers + + block_id = 0 + for index_block, block in enumerate(self.transformer_blocks): + current_text_encoder_layer = text_encoder_layers[block_id] + encoder_hidden_states = torch.cat( + [encoder_hidden_states[:, :, : self.inner_dim // 2], current_text_encoder_layer], dim=-1 + ) + block_id += 1 + if torch.is_grad_enabled() and self.gradient_checkpointing: + encoder_hidden_states, hidden_states = self._gradient_checkpointing_func( + block, + hidden_states, + encoder_hidden_states, + temb, + image_rotary_emb, + joint_attention_kwargs, + ) + + else: + encoder_hidden_states, hidden_states = block( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb=temb, + image_rotary_emb=image_rotary_emb, + joint_attention_kwargs=joint_attention_kwargs, + ) + + for index_block, block in enumerate(self.single_transformer_blocks): + current_text_encoder_layer = text_encoder_layers[block_id] + encoder_hidden_states = torch.cat( + [encoder_hidden_states[:, :, : self.inner_dim // 2], current_text_encoder_layer], dim=-1 + ) + block_id += 1 + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func( + block, + hidden_states, + temb, + image_rotary_emb, + joint_attention_kwargs, + ) + + else: + hidden_states = block( + hidden_states=hidden_states, + temb=temb, + image_rotary_emb=image_rotary_emb, + joint_attention_kwargs=joint_attention_kwargs, + ) + + encoder_hidden_states = hidden_states[:, : encoder_hidden_states.shape[1], ...] + hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...] + + hidden_states = self.norm_out(hidden_states, temb) + output = self.proj_out(hidden_states) + + if USE_PEFT_BACKEND: + # remove `lora_scale` from each PEFT layer + unscale_lora_layers(self, lora_scale) + + if not return_dict: + return (output,) + + return Transformer2DModelOutput(sample=output) diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index 65e8b2469e..db357669b6 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -128,6 +128,7 @@ else: "AnimateDiffVideoToVideoControlNetPipeline", ] _import_structure["bria"] = ["BriaPipeline"] + _import_structure["bria_fibo"] = ["BriaFiboPipeline"] _import_structure["flux"] = [ "FluxControlPipeline", "FluxControlInpaintPipeline", @@ -562,6 +563,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: from .aura_flow import AuraFlowPipeline from .blip_diffusion import BlipDiffusionPipeline from .bria import BriaPipeline + from .bria_fibo import BriaFiboPipeline from .chroma import ChromaImg2ImgPipeline, ChromaPipeline from .cogvideo import ( CogVideoXFunControlPipeline, diff --git a/src/diffusers/pipelines/bria_fibo/__init__.py b/src/diffusers/pipelines/bria_fibo/__init__.py new file mode 100644 index 0000000000..206a463b39 --- /dev/null +++ b/src/diffusers/pipelines/bria_fibo/__init__.py @@ -0,0 +1,48 @@ +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects # noqa F403 + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["pipeline_bria_fibo"] = ["BriaFiboPipeline"] + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() + + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * + else: + from .pipeline_bria_fibo import BriaFiboPipeline + +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) + + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo.py b/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo.py new file mode 100644 index 0000000000..85d29029e6 --- /dev/null +++ b/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo.py @@ -0,0 +1,838 @@ +# Copyright (c) Bria.ai. All rights reserved. +# +# This file is licensed under the Creative Commons Attribution-NonCommercial 4.0 International Public License (CC-BY-NC-4.0). +# You may obtain a copy of the license at https://creativecommons.org/licenses/by-nc/4.0/ +# +# You are free to share and adapt this material for non-commercial purposes provided you give appropriate credit, +# indicate if changes were made, and do not use the material for commercial purposes. +# +# See the license for further details. + +from typing import Any, Callable, Dict, List, Optional, Union + +import numpy as np +import torch +from transformers import AutoTokenizer +from transformers.models.smollm3.modeling_smollm3 import SmolLM3ForCausalLM + +from ...image_processor import VaeImageProcessor +from ...loaders import FluxLoraLoaderMixin +from ...models.autoencoders.autoencoder_kl_wan import AutoencoderKLWan +from ...models.transformers.transformer_bria_fibo import BriaFiboTransformer2DModel +from ...pipelines.bria_fibo.pipeline_output import BriaFiboPipelineOutput +from ...pipelines.flux.pipeline_flux import calculate_shift, retrieve_timesteps +from ...pipelines.pipeline_utils import DiffusionPipeline +from ...schedulers import FlowMatchEulerDiscreteScheduler, KarrasDiffusionSchedulers +from ...utils import ( + USE_PEFT_BACKEND, + is_torch_xla_available, + logging, + replace_example_docstring, + scale_lora_layers, + unscale_lora_layers, +) +from ...utils.torch_utils import randn_tensor + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Example: + ```python + import torch + from diffusers import BriaFiboPipeline + from diffusers.modular_pipelines import ModularPipeline + + torch.set_grad_enabled(False) + vlm_pipe = ModularPipeline.from_pretrained("briaai/FIBO-VLM-prompt-to-JSON", trust_remote_code=True) + + pipe = BriaFiboPipeline.from_pretrained( + "briaai/FIBO", + trust_remote_code=True, + torch_dtype=torch.bfloat16, + ) + pipe.enable_model_cpu_offload() + + with torch.inference_mode(): + # 1. Create a prompt to generate an initial image + output = vlm_pipe(prompt="a beautiful dog") + json_prompt_generate = output.values["json_prompt"] + + # Generate the image from the structured json prompt + results_generate = pipe(prompt=json_prompt_generate, num_inference_steps=50, guidance_scale=5) + results_generate.images[0].save("image_generate.png") + ``` +""" + + +class BriaFiboPipeline(DiffusionPipeline): + r""" + Args: + transformer (`BriaFiboTransformer2DModel`): + The transformer model for 2D diffusion modeling. + scheduler (`FlowMatchEulerDiscreteScheduler` or `KarrasDiffusionSchedulers`): + Scheduler to be used with `transformer` to denoise the encoded latents. + vae (`AutoencoderKLWan`): + Variational Auto-Encoder for encoding and decoding images to and from latent representations. + text_encoder (`SmolLM3ForCausalLM`): + Text encoder for processing input prompts. + tokenizer (`AutoTokenizer`): + Tokenizer used for processing the input text prompts for the text_encoder. + """ + + model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->transformer->vae" + _callback_tensor_inputs = ["latents", "prompt_embeds"] + + def __init__( + self, + transformer: BriaFiboTransformer2DModel, + scheduler: Union[FlowMatchEulerDiscreteScheduler, KarrasDiffusionSchedulers], + vae: AutoencoderKLWan, + text_encoder: SmolLM3ForCausalLM, + tokenizer: AutoTokenizer, + ): + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + scheduler=scheduler, + ) + + self.vae_scale_factor = 16 + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) + self.default_sample_size = 64 + + def get_prompt_embeds( + self, + prompt: Union[str, List[str]], + num_images_per_prompt: int = 1, + max_sequence_length: int = 2048, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + if not prompt: + raise ValueError("`prompt` must be a non-empty string or list of strings.") + + batch_size = len(prompt) + bot_token_id = 128000 + + text_encoder_device = device if device is not None else torch.device("cpu") + if not isinstance(text_encoder_device, torch.device): + text_encoder_device = torch.device(text_encoder_device) + + if all(p == "" for p in prompt): + input_ids = torch.full((batch_size, 1), bot_token_id, dtype=torch.long, device=text_encoder_device) + attention_mask = torch.ones_like(input_ids) + else: + tokenized = self.tokenizer( + prompt, + padding="longest", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + input_ids = tokenized.input_ids.to(text_encoder_device) + attention_mask = tokenized.attention_mask.to(text_encoder_device) + + if any(p == "" for p in prompt): + empty_rows = torch.tensor([p == "" for p in prompt], dtype=torch.bool, device=text_encoder_device) + input_ids[empty_rows] = bot_token_id + attention_mask[empty_rows] = 1 + + encoder_outputs = self.text_encoder( + input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + ) + hidden_states = encoder_outputs.hidden_states + + prompt_embeds = torch.cat([hidden_states[-1], hidden_states[-2]], dim=-1) + prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) + + prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) + hidden_states = tuple( + layer.repeat_interleave(num_images_per_prompt, dim=0).to(device=device) for layer in hidden_states + ) + attention_mask = attention_mask.repeat_interleave(num_images_per_prompt, dim=0).to(device=device) + + return prompt_embeds, hidden_states, attention_mask + + @staticmethod + def pad_embedding(prompt_embeds, max_tokens, attention_mask=None): + # Pad embeddings to `max_tokens` while preserving the mask of real tokens. + batch_size, seq_len, dim = prompt_embeds.shape + + if attention_mask is None: + attention_mask = torch.ones((batch_size, seq_len), dtype=prompt_embeds.dtype, device=prompt_embeds.device) + else: + attention_mask = attention_mask.to(device=prompt_embeds.device, dtype=prompt_embeds.dtype) + + if max_tokens < seq_len: + raise ValueError("`max_tokens` must be greater or equal to the current sequence length.") + + if max_tokens > seq_len: + pad_length = max_tokens - seq_len + padding = torch.zeros( + (batch_size, pad_length, dim), dtype=prompt_embeds.dtype, device=prompt_embeds.device + ) + prompt_embeds = torch.cat([prompt_embeds, padding], dim=1) + + mask_padding = torch.zeros( + (batch_size, pad_length), dtype=prompt_embeds.dtype, device=prompt_embeds.device + ) + attention_mask = torch.cat([attention_mask, mask_padding], dim=1) + + return prompt_embeds, attention_mask + + def encode_prompt( + self, + prompt: Union[str, List[str]], + device: Optional[torch.device] = None, + num_images_per_prompt: int = 1, + guidance_scale: float = 5, + negative_prompt: Optional[Union[str, List[str]]] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + max_sequence_length: int = 3000, + lora_scale: Optional[float] = None, + ): + r""" + Args: + prompt (`str` or `List[str]`, *optional*): + prompt to be encoded + device: (`torch.device`): + torch device + num_images_per_prompt (`int`): + number of images that should be generated per prompt + guidance_scale (`float`): + Guidance scale for classifier free guidance. + 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`). + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + """ + 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, FluxLoraLoaderMixin): + self._lora_scale = lora_scale + + # dynamically adjust the LoRA scale + if self.text_encoder is not None and USE_PEFT_BACKEND: + scale_lora_layers(self.text_encoder, 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] + + prompt_attention_mask = None + negative_prompt_attention_mask = None + if prompt_embeds is None: + prompt_embeds, prompt_layers, prompt_attention_mask = self.get_prompt_embeds( + prompt=prompt, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + ) + prompt_embeds = prompt_embeds.to(dtype=self.transformer.dtype) + prompt_layers = [tensor.to(dtype=self.transformer.dtype) for tensor in prompt_layers] + + if guidance_scale > 1: + if isinstance(negative_prompt, list) and negative_prompt[0] is None: + negative_prompt = "" + negative_prompt = negative_prompt or "" + negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + 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`." + ) + + negative_prompt_embeds, negative_prompt_layers, negative_prompt_attention_mask = self.get_prompt_embeds( + prompt=negative_prompt, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + ) + negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.transformer.dtype) + negative_prompt_layers = [tensor.to(dtype=self.transformer.dtype) for tensor in negative_prompt_layers] + + if self.text_encoder is not None: + if isinstance(self, FluxLoraLoaderMixin) and USE_PEFT_BACKEND: + # Retrieve the original scale by scaling back the LoRA layers + unscale_lora_layers(self.text_encoder, lora_scale) + + # Pad to longest + if prompt_attention_mask is not None: + prompt_attention_mask = prompt_attention_mask.to(device=prompt_embeds.device, dtype=prompt_embeds.dtype) + + if negative_prompt_embeds is not None: + if negative_prompt_attention_mask is not None: + negative_prompt_attention_mask = negative_prompt_attention_mask.to( + device=negative_prompt_embeds.device, dtype=negative_prompt_embeds.dtype + ) + max_tokens = max(negative_prompt_embeds.shape[1], prompt_embeds.shape[1]) + + prompt_embeds, prompt_attention_mask = self.pad_embedding( + prompt_embeds, max_tokens, attention_mask=prompt_attention_mask + ) + prompt_layers = [self.pad_embedding(layer, max_tokens)[0] for layer in prompt_layers] + + negative_prompt_embeds, negative_prompt_attention_mask = self.pad_embedding( + negative_prompt_embeds, max_tokens, attention_mask=negative_prompt_attention_mask + ) + negative_prompt_layers = [self.pad_embedding(layer, max_tokens)[0] for layer in negative_prompt_layers] + else: + max_tokens = prompt_embeds.shape[1] + prompt_embeds, prompt_attention_mask = self.pad_embedding( + prompt_embeds, max_tokens, attention_mask=prompt_attention_mask + ) + negative_prompt_layers = None + + dtype = self.text_encoder.dtype + text_ids = torch.zeros(prompt_embeds.shape[0], max_tokens, 3).to(device=device, dtype=dtype) + + return ( + prompt_embeds, + negative_prompt_embeds, + text_ids, + prompt_attention_mask, + negative_prompt_attention_mask, + prompt_layers, + negative_prompt_layers, + ) + + @property + def guidance_scale(self): + return self._guidance_scale + + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + + @property + def joint_attention_kwargs(self): + return self._joint_attention_kwargs + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def interrupt(self): + return self._interrupt + + @staticmethod + # Based on diffusers.pipelines.flux.pipeline_flux.FluxPipeline._unpack_latents + def _unpack_latents(latents, height, width, vae_scale_factor): + batch_size, num_patches, channels = latents.shape + + height = height // vae_scale_factor + width = width // vae_scale_factor + + latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2) + latents = latents.permute(0, 3, 1, 4, 2, 5) + + latents = latents.reshape(batch_size, channels // (2 * 2), height, width) + return latents + + @staticmethod + # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._prepare_latent_image_ids + def _prepare_latent_image_ids(batch_size, height, width, device, dtype): + latent_image_ids = torch.zeros(height, width, 3) + latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height)[:, None] + latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width)[None, :] + + latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape + + latent_image_ids = latent_image_ids.reshape( + latent_image_id_height * latent_image_id_width, latent_image_id_channels + ) + + return latent_image_ids.to(device=device, dtype=dtype) + + @staticmethod + def _unpack_latents_no_patch(latents, height, width, vae_scale_factor): + batch_size, num_patches, channels = latents.shape + + height = height // vae_scale_factor + width = width // vae_scale_factor + + latents = latents.view(batch_size, height, width, channels) + latents = latents.permute(0, 3, 1, 2) + + return latents + + @staticmethod + def _pack_latents_no_patch(latents, batch_size, num_channels_latents, height, width): + latents = latents.permute(0, 2, 3, 1) + latents = latents.reshape(batch_size, height * width, num_channels_latents) + return latents + + @staticmethod + # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._pack_latents + def _pack_latents(latents, batch_size, num_channels_latents, height, width): + latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2) + latents = latents.permute(0, 2, 4, 1, 3, 5) + latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4) + + return latents + + def prepare_latents( + self, + batch_size, + num_channels_latents, + height, + width, + dtype, + device, + generator, + latents=None, + do_patching=False, + ): + height = int(height) // self.vae_scale_factor + width = int(width) // self.vae_scale_factor + + shape = (batch_size, num_channels_latents, height, width) + + if latents is not None: + latent_image_ids = self._prepare_latent_image_ids(batch_size, height, width, device, dtype) + return latents.to(device=device, dtype=dtype), latent_image_ids + + 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." + ) + + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + if do_patching: + latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width) + latent_image_ids = self._prepare_latent_image_ids(batch_size, height // 2, width // 2, device, dtype) + else: + latents = self._pack_latents_no_patch(latents, batch_size, num_channels_latents, height, width) + latent_image_ids = self._prepare_latent_image_ids(batch_size, height, width, device, dtype) + + return latents, latent_image_ids + + @staticmethod + def _prepare_attention_mask(attention_mask): + attention_matrix = torch.einsum("bi,bj->bij", attention_mask, attention_mask) + + # convert to 0 - keep, -inf ignore + attention_matrix = torch.where( + attention_matrix == 1, 0.0, -torch.inf + ) # Apply -inf to ignored tokens for nulling softmax score + return attention_matrix + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: Union[str, List[str]] = None, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 30, + timesteps: List[int] = None, + guidance_scale: float = 5, + negative_prompt: Optional[Union[str, List[str]]] = None, + num_images_per_prompt: Optional[int] = 1, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + max_sequence_length: int = 3000, + do_patching=False, + ): + r""" + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The height in pixels of the generated image. This is set to 1024 by default for the best results. + width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The width in pixels of the generated image. This is set to 1024 by default for the best results. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + timesteps (`List[int]`, *optional*): + 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. + guidance_scale (`float`, *optional*, defaults to 5.0): + Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). + `guidance_scale` is defined as `w` of equation 2. of [Imagen + Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > + 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, + usually at the expense of lower image quality. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.FloatTensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will ge generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] instead + of a plain tuple. + joint_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + callback_on_step_end (`Callable`, *optional*): + A function that calls at the end of each denoising steps during the inference. The function is called + with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, + callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by + `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + max_sequence_length (`int` defaults to 3000): Maximum sequence length to use with the `prompt`. + do_patching (`bool`, *optional*, defaults to `False`): Whether to use patching. + Examples: + Returns: + [`~pipelines.flux.BriaFiboPipelineOutput`] or `tuple`: [`~pipelines.flux.BriaFiboPipelineOutput`] if + `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the + generated images. + """ + + height = height or self.default_sample_size * self.vae_scale_factor + width = width or self.default_sample_size * self.vae_scale_factor + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt=prompt, + height=height, + width=width, + prompt_embeds=prompt_embeds, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + max_sequence_length=max_sequence_length, + ) + + self._guidance_scale = guidance_scale + self._joint_attention_kwargs = joint_attention_kwargs + self._interrupt = False + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self._execution_device + + lora_scale = ( + self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None + ) + + ( + prompt_embeds, + negative_prompt_embeds, + text_ids, + prompt_attention_mask, + negative_prompt_attention_mask, + prompt_layers, + negative_prompt_layers, + ) = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + guidance_scale=guidance_scale, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + device=device, + max_sequence_length=max_sequence_length, + num_images_per_prompt=num_images_per_prompt, + lora_scale=lora_scale, + ) + prompt_batch_size = prompt_embeds.shape[0] + + if guidance_scale > 1: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + prompt_layers = [ + torch.cat([negative_prompt_layers[i], prompt_layers[i]], dim=0) for i in range(len(prompt_layers)) + ] + prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) + + total_num_layers_transformer = len(self.transformer.transformer_blocks) + len( + self.transformer.single_transformer_blocks + ) + if len(prompt_layers) >= total_num_layers_transformer: + # remove first layers + prompt_layers = prompt_layers[len(prompt_layers) - total_num_layers_transformer :] + else: + # duplicate last layer + prompt_layers = prompt_layers + [prompt_layers[-1]] * (total_num_layers_transformer - len(prompt_layers)) + + # 5. Prepare latent variables + + num_channels_latents = self.transformer.config.in_channels + if do_patching: + num_channels_latents = int(num_channels_latents / 4) + + latents, latent_image_ids = self.prepare_latents( + prompt_batch_size, + num_channels_latents, + height, + width, + prompt_embeds.dtype, + device, + generator, + latents, + do_patching, + ) + + latent_attention_mask = torch.ones( + [latents.shape[0], latents.shape[1]], dtype=latents.dtype, device=latents.device + ) + if guidance_scale > 1: + latent_attention_mask = latent_attention_mask.repeat(2, 1) + + attention_mask = torch.cat([prompt_attention_mask, latent_attention_mask], dim=1) + attention_mask = self._prepare_attention_mask(attention_mask) # batch, seq => batch, seq, seq + attention_mask = attention_mask.unsqueeze(dim=1).to(dtype=self.transformer.dtype) # for head broadcasting + + if self._joint_attention_kwargs is None: + self._joint_attention_kwargs = {} + self._joint_attention_kwargs["attention_mask"] = attention_mask + + # Adapt scheduler to dynamic shifting (resolution dependent) + + if do_patching: + seq_len = (height // (self.vae_scale_factor * 2)) * (width // (self.vae_scale_factor * 2)) + else: + seq_len = (height // self.vae_scale_factor) * (width // self.vae_scale_factor) + + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) + + mu = calculate_shift( + seq_len, + self.scheduler.config.base_image_seq_len, + self.scheduler.config.max_image_seq_len, + self.scheduler.config.base_shift, + self.scheduler.config.max_shift, + ) + + # Init sigmas and timesteps according to shift size + # This changes the scheduler in-place according to the dynamic scheduling + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps=num_inference_steps, + device=device, + timesteps=None, + sigmas=sigmas, + mu=mu, + ) + + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + + # Support old different diffusers versions + if len(latent_image_ids.shape) == 3: + latent_image_ids = latent_image_ids[0] + + if len(text_ids.shape) == 3: + text_ids = text_ids[0] + + # 6. Denoising loop + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat([latents] * 2) if guidance_scale > 1 else latents + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latent_model_input.shape[0]).to( + device=latent_model_input.device, dtype=latent_model_input.dtype + ) + + # This is predicts "v" from flow-matching or eps from diffusion + noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=prompt_embeds, + text_encoder_layers=prompt_layers, + joint_attention_kwargs=self.joint_attention_kwargs, + return_dict=False, + txt_ids=text_ids, + img_ids=latent_image_ids, + )[0] + + # perform guidance + if guidance_scale > 1: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents_dtype = latents.dtype + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + + if latents.dtype != latents_dtype: + if torch.backends.mps.is_available(): + # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272 + latents = latents.to(latents_dtype) + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + if output_type == "latent": + image = latents + + else: + if do_patching: + latents = self._unpack_latents(latents, height, width, self.vae_scale_factor) + else: + latents = self._unpack_latents_no_patch(latents, height, width, self.vae_scale_factor) + + latents = latents.unsqueeze(dim=2) + latents_device = latents[0].device + latents_dtype = latents[0].dtype + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(latents_device, latents_dtype) + ) + latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( + latents_device, latents_dtype + ) + latents_scaled = [latent / latents_std + latents_mean for latent in latents] + latents_scaled = torch.cat(latents_scaled, dim=0) + image = [] + for scaled_latent in latents_scaled: + curr_image = self.vae.decode(scaled_latent.unsqueeze(0), return_dict=False)[0] + curr_image = self.image_processor.postprocess(curr_image.squeeze(dim=2), output_type=output_type) + image.append(curr_image) + if len(image) == 1: + image = image[0] + else: + image = np.stack(image, axis=0) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return BriaFiboPipelineOutput(images=image) + + def check_inputs( + self, + prompt, + height, + width, + negative_prompt=None, + prompt_embeds=None, + negative_prompt_embeds=None, + callback_on_step_end_tensor_inputs=None, + max_sequence_length=None, + ): + if height % 16 != 0 or width % 16 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.") + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + if prompt_embeds is not None and negative_prompt_embeds is not None: + if prompt_embeds.shape != negative_prompt_embeds.shape: + raise ValueError( + "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" + f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" + f" {negative_prompt_embeds.shape}." + ) + + if max_sequence_length is not None and max_sequence_length > 3000: + raise ValueError(f"`max_sequence_length` cannot be greater than 3000 but is {max_sequence_length}") diff --git a/src/diffusers/pipelines/bria_fibo/pipeline_output.py b/src/diffusers/pipelines/bria_fibo/pipeline_output.py new file mode 100644 index 0000000000..f459185a2c --- /dev/null +++ b/src/diffusers/pipelines/bria_fibo/pipeline_output.py @@ -0,0 +1,21 @@ +from dataclasses import dataclass +from typing import List, Union + +import numpy as np +import PIL.Image + +from ...utils import BaseOutput + + +@dataclass +class BriaFiboPipelineOutput(BaseOutput): + """ + Output class for BriaFibo pipelines. + + Args: + images (`List[PIL.Image.Image]` or `np.ndarray`) + List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, + num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. + """ + + images: Union[List[PIL.Image.Image], np.ndarray] diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 6ef6b8b0e9..3c426d5039 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -588,6 +588,21 @@ class AutoModel(metaclass=DummyObject): requires_backends(cls, ["torch"]) +class BriaFiboTransformer2DModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class BriaTransformer2DModel(metaclass=DummyObject): _backends = ["torch"] diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 7c1dcba9c7..20575ff229 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -482,6 +482,21 @@ class AuraFlowPipeline(metaclass=DummyObject): requires_backends(cls, ["torch", "transformers"]) +class BriaFiboPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class BriaPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/tests/models/transformers/test_models_transformer_bria_fibo.py b/tests/models/transformers/test_models_transformer_bria_fibo.py new file mode 100644 index 0000000000..f859f4608b --- /dev/null +++ b/tests/models/transformers/test_models_transformer_bria_fibo.py @@ -0,0 +1,89 @@ +# coding=utf-8 +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import torch + +from diffusers import BriaFiboTransformer2DModel + +from ...testing_utils import enable_full_determinism, torch_device +from ..test_modeling_common import ModelTesterMixin + + +enable_full_determinism() + + +class BriaFiboTransformerTests(ModelTesterMixin, unittest.TestCase): + model_class = BriaFiboTransformer2DModel + main_input_name = "hidden_states" + # We override the items here because the transformer under consideration is small. + model_split_percents = [0.8, 0.7, 0.7] + + # Skip setting testing with default: AttnProcessor + uses_custom_attn_processor = True + + @property + def dummy_input(self): + batch_size = 1 + num_latent_channels = 48 + num_image_channels = 3 + height = width = 16 + sequence_length = 32 + embedding_dim = 64 + + hidden_states = torch.randn((batch_size, height * width, num_latent_channels)).to(torch_device) + encoder_hidden_states = torch.randn((batch_size, sequence_length, embedding_dim)).to(torch_device) + text_ids = torch.randn((sequence_length, num_image_channels)).to(torch_device) + image_ids = torch.randn((height * width, num_image_channels)).to(torch_device) + timestep = torch.tensor([1.0]).to(torch_device).expand(batch_size) + + return { + "hidden_states": hidden_states, + "encoder_hidden_states": encoder_hidden_states, + "img_ids": image_ids, + "txt_ids": text_ids, + "timestep": timestep, + "text_encoder_layers": [encoder_hidden_states[:, :, :32], encoder_hidden_states[:, :, :32]], + } + + @property + def input_shape(self): + return (16, 16) + + @property + def output_shape(self): + return (256, 48) + + def prepare_init_args_and_inputs_for_common(self): + init_dict = { + "patch_size": 1, + "in_channels": 48, + "num_layers": 1, + "num_single_layers": 1, + "attention_head_dim": 8, + "num_attention_heads": 2, + "joint_attention_dim": 64, + "text_encoder_dim": 32, + "pooled_projection_dim": None, + "axes_dims_rope": [0, 4, 4], + } + + inputs_dict = self.dummy_input + return init_dict, inputs_dict + + def test_gradient_checkpointing_is_applied(self): + expected_set = {"BriaFiboTransformer2DModel"} + super().test_gradient_checkpointing_is_applied(expected_set=expected_set) diff --git a/tests/pipelines/bria_fibo/__init__.py b/tests/pipelines/bria_fibo/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/pipelines/bria_fibo/test_pipeline_bria_fibo.py b/tests/pipelines/bria_fibo/test_pipeline_bria_fibo.py new file mode 100644 index 0000000000..76b41114f8 --- /dev/null +++ b/tests/pipelines/bria_fibo/test_pipeline_bria_fibo.py @@ -0,0 +1,139 @@ +# Copyright 2024 Bria AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import numpy as np +import torch +from transformers import AutoTokenizer +from transformers.models.smollm3.modeling_smollm3 import SmolLM3Config, SmolLM3ForCausalLM + +from diffusers import ( + AutoencoderKLWan, + BriaFiboPipeline, + FlowMatchEulerDiscreteScheduler, +) +from diffusers.models.transformers.transformer_bria_fibo import BriaFiboTransformer2DModel +from tests.pipelines.test_pipelines_common import PipelineTesterMixin + +from ...testing_utils import ( + enable_full_determinism, + torch_device, +) + + +enable_full_determinism() + + +class BriaFiboPipelineFastTests(PipelineTesterMixin, unittest.TestCase): + pipeline_class = BriaFiboPipeline + params = frozenset(["prompt", "height", "width", "guidance_scale"]) + batch_params = frozenset(["prompt"]) + test_xformers_attention = False + test_layerwise_casting = False + test_group_offloading = False + supports_dduf = False + + def get_dummy_components(self): + torch.manual_seed(0) + transformer = BriaFiboTransformer2DModel( + patch_size=1, + in_channels=16, + num_layers=1, + num_single_layers=1, + attention_head_dim=8, + num_attention_heads=2, + joint_attention_dim=64, + text_encoder_dim=32, + pooled_projection_dim=None, + axes_dims_rope=[0, 4, 4], + ) + + torch.manual_seed(0) + vae = AutoencoderKLWan( + base_dim=160, + decoder_base_dim=256, + num_res_blocks=2, + out_channels=12, + patch_size=2, + scale_factor_spatial=16, + scale_factor_temporal=4, + temperal_downsample=[False, True, True], + z_dim=16, + ) + + scheduler = FlowMatchEulerDiscreteScheduler() + + torch.manual_seed(0) + text_encoder = SmolLM3ForCausalLM(SmolLM3Config(hidden_size=32)) + tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") + + components = { + "scheduler": scheduler, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + "transformer": transformer, + "vae": vae, + } + return components + + def get_dummy_inputs(self, device, seed=0): + if str(device).startswith("mps"): + generator = torch.manual_seed(seed) + else: + generator = torch.Generator(device="cpu").manual_seed(seed) + + inputs = { + "prompt": "{'text': 'A painting of a squirrel eating a burger'}", + "negative_prompt": "bad, ugly", + "generator": generator, + "num_inference_steps": 2, + "guidance_scale": 5.0, + "height": 32, + "width": 32, + "output_type": "np", + } + return inputs + + @unittest.skip(reason="will not be supported due to dim-fusion") + def test_encode_prompt_works_in_isolation(self): + pass + + def test_bria_fibo_different_prompts(self): + pipe = self.pipeline_class(**self.get_dummy_components()) + pipe = pipe.to(torch_device) + inputs = self.get_dummy_inputs(torch_device) + output_same_prompt = pipe(**inputs).images[0] + + inputs = self.get_dummy_inputs(torch_device) + inputs["prompt"] = "a different prompt" + output_different_prompts = pipe(**inputs).images[0] + + max_diff = np.abs(output_same_prompt - output_different_prompts).max() + assert max_diff > 1e-6 + + def test_image_output_shape(self): + pipe = self.pipeline_class(**self.get_dummy_components()) + pipe = pipe.to(torch_device) + inputs = self.get_dummy_inputs(torch_device) + + height_width_pairs = [(32, 32), (64, 64), (32, 64)] + for height, width in height_width_pairs: + expected_height = height + expected_width = width + + inputs.update({"height": height, "width": width}) + image = pipe(**inputs).images[0] + output_height, output_width, _ = image.shape + assert (output_height, output_width) == (expected_height, expected_width) From 9f3c0fdcd859905c2c13ec47f10eb0250d2576ac Mon Sep 17 00:00:00 2001 From: Pavle Padjin Date: Thu, 30 Oct 2025 04:09:40 +0100 Subject: [PATCH 107/123] Avoiding graph break by changing the way we infer dtype in vae.decoder (#12512) * Changing the way we infer dtype to avoid force evaluation of lazy tensors * changing way to infer dtype to ensure type consistency * more robust infering of dtype * removing the upscale dtype entirely --- src/diffusers/models/autoencoders/vae.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/diffusers/models/autoencoders/vae.py b/src/diffusers/models/autoencoders/vae.py index 7b17196125..9c6031a988 100644 --- a/src/diffusers/models/autoencoders/vae.py +++ b/src/diffusers/models/autoencoders/vae.py @@ -286,11 +286,9 @@ class Decoder(nn.Module): sample = self.conv_in(sample) - upscale_dtype = next(iter(self.up_blocks.parameters())).dtype if torch.is_grad_enabled() and self.gradient_checkpointing: # middle sample = self._gradient_checkpointing_func(self.mid_block, sample, latent_embeds) - sample = sample.to(upscale_dtype) # up for up_block in self.up_blocks: @@ -298,7 +296,6 @@ class Decoder(nn.Module): else: # middle sample = self.mid_block(sample, latent_embeds) - sample = sample.to(upscale_dtype) # up for up_block in self.up_blocks: From df8dd778177c3d2272f74cbebd880d7abd9f5ec9 Mon Sep 17 00:00:00 2001 From: Dhruv Nair Date: Fri, 31 Oct 2025 00:14:24 +0530 Subject: [PATCH 108/123] [Modular] Fix for custom block kwargs (#12561) update --- src/diffusers/modular_pipelines/modular_pipeline.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index bf067555a8..55c261ab2f 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -335,7 +335,7 @@ class ModularPipelineBlocks(ConfigMixin, PushToHubMixin): ) expected_kwargs, optional_kwargs = block_cls._get_signature_keys(block_cls) block_kwargs = { - name: kwargs.pop(name) for name in kwargs if name in expected_kwargs or name in optional_kwargs + name: kwargs.get(name) for name in kwargs if name in expected_kwargs or name in optional_kwargs } return block_cls(**block_kwargs) From d54622c2679d700b425ad61abce9b80fc36212c0 Mon Sep 17 00:00:00 2001 From: Dhruv Nair Date: Fri, 31 Oct 2025 13:47:02 +0530 Subject: [PATCH 109/123] [Modular] Allow custom blocks to be saved to `local_dir` (#12381) update Co-authored-by: YiYi Xu --- src/diffusers/modular_pipelines/modular_pipeline.py | 2 +- src/diffusers/utils/dynamic_modules_utils.py | 7 ++++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 55c261ab2f..ef1673c057 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -305,6 +305,7 @@ class ModularPipelineBlocks(ConfigMixin, PushToHubMixin): "cache_dir", "force_download", "local_files_only", + "local_dir", "proxies", "resume_download", "revision", @@ -331,7 +332,6 @@ class ModularPipelineBlocks(ConfigMixin, PushToHubMixin): module_file=module_file, class_name=class_name, **hub_kwargs, - **kwargs, ) expected_kwargs, optional_kwargs = block_cls._get_signature_keys(block_cls) block_kwargs = { diff --git a/src/diffusers/utils/dynamic_modules_utils.py b/src/diffusers/utils/dynamic_modules_utils.py index 627b1e0604..b2ef5a29e0 100644 --- a/src/diffusers/utils/dynamic_modules_utils.py +++ b/src/diffusers/utils/dynamic_modules_utils.py @@ -254,6 +254,7 @@ def get_cached_module_file( token: Optional[Union[bool, str]] = None, revision: Optional[str] = None, local_files_only: bool = False, + local_dir: Optional[str] = None, ): """ Prepares Downloads a module from a local folder or a distant repo and returns its path inside the cached @@ -332,6 +333,7 @@ def get_cached_module_file( force_download=force_download, proxies=proxies, local_files_only=local_files_only, + local_dir=local_dir, ) submodule = "git" module_file = pretrained_model_name_or_path + ".py" @@ -355,6 +357,7 @@ def get_cached_module_file( force_download=force_download, proxies=proxies, local_files_only=local_files_only, + local_dir=local_dir, token=token, ) submodule = os.path.join("local", "--".join(pretrained_model_name_or_path.split("/"))) @@ -415,6 +418,7 @@ def get_cached_module_file( token=token, revision=revision, local_files_only=local_files_only, + local_dir=local_dir, ) return os.path.join(full_submodule, module_file) @@ -431,7 +435,7 @@ def get_class_from_dynamic_module( token: Optional[Union[bool, str]] = None, revision: Optional[str] = None, local_files_only: bool = False, - **kwargs, + local_dir: Optional[str] = None, ): """ Extracts a class from a module file, present in the local folder or repository of a model. @@ -496,5 +500,6 @@ def get_class_from_dynamic_module( token=token, revision=revision, local_files_only=local_files_only, + local_dir=local_dir, ) return get_class_in_module(class_name, final_module) From 051c8a1c0f5c393a447bef18081fdf94c2a3ab9e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Friedrich=20Sch=C3=B6ller?= Date: Fri, 31 Oct 2025 21:25:13 +0100 Subject: [PATCH 110/123] Fix Stable Diffusion 3.x pooled prompt embedding with multiple images (#12306) --- .../controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py | 2 +- .../pipeline_stable_diffusion_3_controlnet_inpainting.py | 2 +- src/diffusers/pipelines/pag/pipeline_pag_sd_3.py | 2 +- src/diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py | 2 +- .../pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py | 2 +- .../stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py | 2 +- .../stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py | 2 +- 7 files changed, 7 insertions(+), 7 deletions(-) diff --git a/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py b/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py index f67a0e2112..d605eac1f2 100644 --- a/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py +++ b/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py @@ -355,7 +355,7 @@ class StableDiffusion3ControlNetPipeline( prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) - pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt, 1) + pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt) pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) return prompt_embeds, pooled_prompt_embeds diff --git a/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py b/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py index 68984da4dc..9d0158c6b6 100644 --- a/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py +++ b/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py @@ -373,7 +373,7 @@ class StableDiffusion3ControlNetInpaintingPipeline( prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) - pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt, 1) + pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt) pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) return prompt_embeds, pooled_prompt_embeds diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py index bc281428e2..941b675099 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py @@ -326,7 +326,7 @@ class StableDiffusion3PAGPipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSin prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) - pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt, 1) + pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt) pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) return prompt_embeds, pooled_prompt_embeds diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py index 22a8dac238..f40dd52fc2 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py @@ -342,7 +342,7 @@ class StableDiffusion3PAGImg2ImgPipeline(DiffusionPipeline, SD3LoraLoaderMixin, prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) - pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt, 1) + pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt) pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) return prompt_embeds, pooled_prompt_embeds diff --git a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py index 3b7b26dc63..660d9801df 100644 --- a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py +++ b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py @@ -336,7 +336,7 @@ class StableDiffusion3Pipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingle prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) - pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt, 1) + pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt) pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) return prompt_embeds, pooled_prompt_embeds diff --git a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py index db047f1992..9b11bc8781 100644 --- a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py +++ b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py @@ -361,7 +361,7 @@ class StableDiffusion3Img2ImgPipeline(DiffusionPipeline, SD3LoraLoaderMixin, Fro prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) - pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt, 1) + pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt) pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) return prompt_embeds, pooled_prompt_embeds diff --git a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py index c95fa530c8..b947cbff09 100644 --- a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py +++ b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py @@ -367,7 +367,7 @@ class StableDiffusion3InpaintPipeline(DiffusionPipeline, SD3LoraLoaderMixin, Fro prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) - pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt, 1) + pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt) pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) return prompt_embeds, pooled_prompt_embeds From 5e8415a3116b045798c973fe6dce0c9efc130d03 Mon Sep 17 00:00:00 2001 From: Dhruv Nair Date: Sun, 2 Nov 2025 08:34:31 +0530 Subject: [PATCH 111/123] Fix custom code loading in Automodel (#12571) update --- src/diffusers/models/auto_model.py | 4 +--- src/diffusers/modular_pipelines/modular_pipeline.py | 1 - 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/src/diffusers/models/auto_model.py b/src/diffusers/models/auto_model.py index 947b610ea6..c96b4fa88c 100644 --- a/src/diffusers/models/auto_model.py +++ b/src/diffusers/models/auto_model.py @@ -147,14 +147,13 @@ class AutoModel(ConfigMixin): "force_download", "local_files_only", "proxies", - "resume_download", "revision", "token", ] hub_kwargs = {name: kwargs.pop(name, None) for name in hub_kwargs_names} # load_config_kwargs uses the same hub kwargs minus subfolder and resume_download - load_config_kwargs = {k: v for k, v in hub_kwargs.items() if k not in ["subfolder", "resume_download"]} + load_config_kwargs = {k: v for k, v in hub_kwargs.items() if k not in ["subfolder"]} library = None orig_class_name = None @@ -205,7 +204,6 @@ class AutoModel(ConfigMixin): module_file=module_file, class_name=class_name, **hub_kwargs, - **kwargs, ) else: from ..pipelines.pipeline_loading_utils import ALL_IMPORTABLE_CLASSES, get_class_obj_and_candidates diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index ef1673c057..e6ec32ed9c 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -307,7 +307,6 @@ class ModularPipelineBlocks(ConfigMixin, PushToHubMixin): "local_files_only", "local_dir", "proxies", - "resume_download", "revision", "subfolder", "token", From cdbf0ad8839d6e3b4111a87fddc55ff798d2d8a7 Mon Sep 17 00:00:00 2001 From: YiYi Xu Date: Sat, 1 Nov 2025 18:45:09 -1000 Subject: [PATCH 112/123] [modular] better warn message (#12573) better warn message --- src/diffusers/modular_pipelines/modular_pipeline.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index e6ec32ed9c..8076f41c15 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -2130,8 +2130,13 @@ class ModularPipeline(ConfigMixin, PushToHubMixin): component_load_kwargs[key] = value["default"] try: components_to_register[name] = spec.load(**component_load_kwargs) - except Exception as e: - logger.warning(f"Failed to create component '{name}': {e}") + except Exception: + logger.warning( + f"\nFailed to create component {name}:\n" + f"- Component spec: {spec}\n" + f"- load() called with kwargs: {component_load_kwargs}\n\n" + f"{traceback.format_exc()}" + ) # Register all components at once self.register_components(**components_to_register) From 8f80dda193f79af3ccd0f985906d61123d69df08 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Sun, 2 Nov 2025 10:51:11 +0530 Subject: [PATCH 113/123] [tests] add tests for flux modular (t2i, i2i, kontext) (#12566) * start flux modular tests. * up * add kontext * up * up * up * Update src/diffusers/modular_pipelines/flux/denoise.py Co-authored-by: YiYi Xu * up * up --------- Co-authored-by: YiYi Xu --- .../modular_pipelines/components_manager.py | 8 +- .../modular_pipelines/flux/before_denoise.py | 2 +- .../modular_pipelines/flux/denoise.py | 4 +- .../modular_pipelines/flux/encoders.py | 11 +- .../modular_pipelines/flux/inputs.py | 4 + tests/modular_pipelines/flux/__init__.py | 0 .../flux/test_modular_pipeline_flux.py | 130 ++++++++++++++++++ ...st_modular_pipeline_stable_diffusion_xl.py | 20 +-- 8 files changed, 152 insertions(+), 27 deletions(-) create mode 100644 tests/modular_pipelines/flux/__init__.py create mode 100644 tests/modular_pipelines/flux/test_modular_pipeline_flux.py diff --git a/src/diffusers/modular_pipelines/components_manager.py b/src/diffusers/modular_pipelines/components_manager.py index 9dd8035c44..cb7e8fb736 100644 --- a/src/diffusers/modular_pipelines/components_manager.py +++ b/src/diffusers/modular_pipelines/components_manager.py @@ -164,7 +164,11 @@ class AutoOffloadStrategy: device_type = execution_device.type device_module = getattr(torch, device_type, torch.cuda) - mem_on_device = device_module.mem_get_info(execution_device.index)[0] + try: + mem_on_device = device_module.mem_get_info(execution_device.index)[0] + except AttributeError: + raise AttributeError(f"Do not know how to obtain obtain memory info for {str(device_module)}.") + mem_on_device = mem_on_device - self.memory_reserve_margin if current_module_size < mem_on_device: return [] @@ -699,6 +703,8 @@ class ComponentsManager: if not is_accelerate_available(): raise ImportError("Make sure to install accelerate to use auto_cpu_offload") + # TODO: add a warning if mem_get_info isn't available on `device`. + for name, component in self.components.items(): if isinstance(component, torch.nn.Module) and hasattr(component, "_hf_hook"): remove_hook_from_module(component, recurse=True) diff --git a/src/diffusers/modular_pipelines/flux/before_denoise.py b/src/diffusers/modular_pipelines/flux/before_denoise.py index c098b7d4f1..daffec9865 100644 --- a/src/diffusers/modular_pipelines/flux/before_denoise.py +++ b/src/diffusers/modular_pipelines/flux/before_denoise.py @@ -598,7 +598,7 @@ class FluxKontextRoPEInputsStep(ModularPipelineBlocks): and getattr(block_state, "image_width", None) is not None ): image_latent_height = 2 * (int(block_state.image_height) // (components.vae_scale_factor * 2)) - image_latent_width = 2 * (int(block_state.width) // (components.vae_scale_factor * 2)) + image_latent_width = 2 * (int(block_state.image_width) // (components.vae_scale_factor * 2)) img_ids = FluxPipeline._prepare_latent_image_ids( None, image_latent_height // 2, image_latent_width // 2, device, dtype ) diff --git a/src/diffusers/modular_pipelines/flux/denoise.py b/src/diffusers/modular_pipelines/flux/denoise.py index b1796bb63c..5a769df103 100644 --- a/src/diffusers/modular_pipelines/flux/denoise.py +++ b/src/diffusers/modular_pipelines/flux/denoise.py @@ -59,7 +59,7 @@ class FluxLoopDenoiser(ModularPipelineBlocks): ), InputParam( "guidance", - required=True, + required=False, type_hint=torch.Tensor, description="Guidance scale as a tensor", ), @@ -141,7 +141,7 @@ class FluxKontextLoopDenoiser(ModularPipelineBlocks): ), InputParam( "guidance", - required=True, + required=False, type_hint=torch.Tensor, description="Guidance scale as a tensor", ), diff --git a/src/diffusers/modular_pipelines/flux/encoders.py b/src/diffusers/modular_pipelines/flux/encoders.py index b71962bd93..f0314d4771 100644 --- a/src/diffusers/modular_pipelines/flux/encoders.py +++ b/src/diffusers/modular_pipelines/flux/encoders.py @@ -95,7 +95,7 @@ class FluxProcessImagesInputStep(ModularPipelineBlocks): ComponentSpec( "image_processor", VaeImageProcessor, - config=FrozenDict({"vae_scale_factor": 16}), + config=FrozenDict({"vae_scale_factor": 16, "vae_latent_channels": 16}), default_creation_method="from_config", ), ] @@ -143,10 +143,6 @@ class FluxProcessImagesInputStep(ModularPipelineBlocks): class FluxKontextProcessImagesInputStep(ModularPipelineBlocks): model_name = "flux-kontext" - def __init__(self, _auto_resize=True): - self._auto_resize = _auto_resize - super().__init__() - @property def description(self) -> str: return ( @@ -167,7 +163,7 @@ class FluxKontextProcessImagesInputStep(ModularPipelineBlocks): @property def inputs(self) -> List[InputParam]: - return [InputParam("image")] + return [InputParam("image"), InputParam("_auto_resize", type_hint=bool, default=True)] @property def intermediate_outputs(self) -> List[OutputParam]: @@ -195,7 +191,8 @@ class FluxKontextProcessImagesInputStep(ModularPipelineBlocks): img = images[0] image_height, image_width = components.image_processor.get_default_height_width(img) aspect_ratio = image_width / image_height - if self._auto_resize: + _auto_resize = block_state._auto_resize + if _auto_resize: # Kontext is trained on specific resolutions, using one of them is recommended _, image_width, image_height = min( (abs(aspect_ratio - w / h), w, h) for w, h in PREFERRED_KONTEXT_RESOLUTIONS diff --git a/src/diffusers/modular_pipelines/flux/inputs.py b/src/diffusers/modular_pipelines/flux/inputs.py index e1bc17f5ff..8309eebfeb 100644 --- a/src/diffusers/modular_pipelines/flux/inputs.py +++ b/src/diffusers/modular_pipelines/flux/inputs.py @@ -112,6 +112,10 @@ class FluxTextInputStep(ModularPipelineBlocks): block_state.prompt_embeds = block_state.prompt_embeds.view( block_state.batch_size * block_state.num_images_per_prompt, seq_len, -1 ) + pooled_prompt_embeds = block_state.pooled_prompt_embeds.repeat(1, block_state.num_images_per_prompt) + block_state.pooled_prompt_embeds = pooled_prompt_embeds.view( + block_state.batch_size * block_state.num_images_per_prompt, -1 + ) self.set_block_state(state, block_state) return components, state diff --git a/tests/modular_pipelines/flux/__init__.py b/tests/modular_pipelines/flux/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py new file mode 100644 index 0000000000..9d70c21aa8 --- /dev/null +++ b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py @@ -0,0 +1,130 @@ +# coding=utf-8 +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import random +import tempfile +import unittest + +import numpy as np +import PIL +import torch + +from diffusers.image_processor import VaeImageProcessor +from diffusers.modular_pipelines import ( + FluxAutoBlocks, + FluxKontextAutoBlocks, + FluxKontextModularPipeline, + FluxModularPipeline, + ModularPipeline, +) + +from ...testing_utils import floats_tensor, torch_device +from ..test_modular_pipelines_common import ModularPipelineTesterMixin + + +class FluxModularTests: + pipeline_class = FluxModularPipeline + pipeline_blocks_class = FluxAutoBlocks + repo = "hf-internal-testing/tiny-flux-modular" + + def get_pipeline(self, components_manager=None, torch_dtype=torch.float32): + pipeline = self.pipeline_blocks_class().init_pipeline(self.repo, components_manager=components_manager) + pipeline.load_components(torch_dtype=torch_dtype) + return pipeline + + def get_dummy_inputs(self, device, seed=0): + if str(device).startswith("mps"): + generator = torch.manual_seed(seed) + else: + generator = torch.Generator(device=device).manual_seed(seed) + inputs = { + "prompt": "A painting of a squirrel eating a burger", + "generator": generator, + "num_inference_steps": 2, + "guidance_scale": 5.0, + "height": 8, + "width": 8, + "max_sequence_length": 48, + "output_type": "np", + } + return inputs + + +class FluxModularPipelineFastTests(FluxModularTests, ModularPipelineTesterMixin, unittest.TestCase): + params = frozenset(["prompt", "height", "width", "guidance_scale"]) + batch_params = frozenset(["prompt"]) + + +class FluxImg2ImgModularPipelineFastTests(FluxModularTests, ModularPipelineTesterMixin, unittest.TestCase): + params = frozenset(["prompt", "height", "width", "guidance_scale", "image"]) + batch_params = frozenset(["prompt", "image"]) + + def get_pipeline(self, components_manager=None, torch_dtype=torch.float32): + pipeline = super().get_pipeline(components_manager, torch_dtype) + # Override `vae_scale_factor` here as currently, `image_processor` is initialized with + # fixed constants instead of + # https://github.com/huggingface/diffusers/blob/d54622c2679d700b425ad61abce9b80fc36212c0/src/diffusers/pipelines/flux/pipeline_flux_img2img.py#L230C9-L232C10 + pipeline.image_processor = VaeImageProcessor(vae_scale_factor=2) + return pipeline + + def get_dummy_inputs(self, device, seed=0): + inputs = super().get_dummy_inputs(device, seed) + image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device) + image = image / 2 + 0.5 + inputs["image"] = image + inputs["strength"] = 0.8 + inputs["height"] = 8 + inputs["width"] = 8 + return inputs + + def test_save_from_pretrained(self): + pipes = [] + base_pipe = self.get_pipeline().to(torch_device) + pipes.append(base_pipe) + + with tempfile.TemporaryDirectory() as tmpdirname: + base_pipe.save_pretrained(tmpdirname) + pipe = ModularPipeline.from_pretrained(tmpdirname).to(torch_device) + pipe.load_components(torch_dtype=torch.float32) + pipe.to(torch_device) + pipe.image_processor = VaeImageProcessor(vae_scale_factor=2) + + pipes.append(pipe) + + image_slices = [] + for pipe in pipes: + inputs = self.get_dummy_inputs(torch_device) + image = pipe(**inputs, output="images") + + image_slices.append(image[0, -3:, -3:, -1].flatten()) + + assert np.abs(image_slices[0] - image_slices[1]).max() < 1e-3 + + +class FluxKontextModularPipelineFastTests(FluxImg2ImgModularPipelineFastTests): + pipeline_class = FluxKontextModularPipeline + pipeline_blocks_class = FluxKontextAutoBlocks + repo = "hf-internal-testing/tiny-flux-kontext-pipe" + + def get_dummy_inputs(self, device, seed=0): + inputs = super().get_dummy_inputs(device, seed) + image = PIL.Image.new("RGB", (32, 32), 0) + _ = inputs.pop("strength") + inputs["image"] = image + inputs["height"] = 8 + inputs["width"] = 8 + inputs["max_area"] = 8 * 8 + inputs["_auto_resize"] = False + return inputs diff --git a/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py b/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py index d05f818135..22347aa558 100644 --- a/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py +++ b/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py @@ -21,24 +21,12 @@ import numpy as np import torch from PIL import Image -from diffusers import ( - ClassifierFreeGuidance, - StableDiffusionXLAutoBlocks, - StableDiffusionXLModularPipeline, -) +from diffusers import ClassifierFreeGuidance, StableDiffusionXLAutoBlocks, StableDiffusionXLModularPipeline from diffusers.loaders import ModularIPAdapterMixin -from ...models.unets.test_models_unet_2d_condition import ( - create_ip_adapter_state_dict, -) -from ...testing_utils import ( - enable_full_determinism, - floats_tensor, - torch_device, -) -from ..test_modular_pipelines_common import ( - ModularPipelineTesterMixin, -) +from ...models.unets.test_models_unet_2d_condition import create_ip_adapter_state_dict +from ...testing_utils import enable_full_determinism, floats_tensor, torch_device +from ..test_modular_pipelines_common import ModularPipelineTesterMixin enable_full_determinism() From de6173c683d72feadc039d30d323427d6265e616 Mon Sep 17 00:00:00 2001 From: YiYi Xu Date: Mon, 3 Nov 2025 09:44:42 -1000 Subject: [PATCH 114/123] [modular]pass hub_kwargs to load_config (#12577) pass hub_kwargs to load_config --- src/diffusers/modular_pipelines/modular_pipeline.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 8076f41c15..307698245e 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -313,7 +313,7 @@ class ModularPipelineBlocks(ConfigMixin, PushToHubMixin): ] hub_kwargs = {name: kwargs.pop(name) for name in hub_kwargs_names if name in kwargs} - config = cls.load_config(pretrained_model_name_or_path) + config = cls.load_config(pretrained_model_name_or_path, **hub_kwargs) has_remote_code = "auto_map" in config and cls.__name__ in config["auto_map"] trust_remote_code = resolve_trust_remote_code( trust_remote_code, pretrained_model_name_or_path, has_remote_code From 1ec28a2c770546e4483a5109c4e98d6e1b298968 Mon Sep 17 00:00:00 2001 From: "Wang, Yi" Date: Tue, 4 Nov 2025 05:48:20 +0800 Subject: [PATCH 115/123] ulysses enabling in native attention path (#12563) * ulysses enabling in native attention path Signed-off-by: Wang, Yi A * address review comment Signed-off-by: Wang, Yi A * add supports_context_parallel for native attention Signed-off-by: Wang, Yi A * update templated attention Signed-off-by: Wang, Yi A --------- Signed-off-by: Wang, Yi A Co-authored-by: Sayak Paul --- src/diffusers/models/attention_dispatch.py | 122 +++++++++++++++++++-- 1 file changed, 110 insertions(+), 12 deletions(-) diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index ab0d7102ee..c17a3d0ed6 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -649,6 +649,86 @@ def _( # ===== Helper functions to use attention backends with templated CP autograd functions ===== +def _native_attention_forward_op( + ctx: torch.autograd.function.FunctionCtx, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: Optional[torch.Tensor] = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: Optional[float] = None, + enable_gqa: bool = False, + return_lse: bool = False, + _save_ctx: bool = True, + _parallel_config: Optional["ParallelConfig"] = None, +): + # Native attention does not return_lse + if return_lse: + raise ValueError("Native attention does not support return_lse=True") + + # used for backward pass + if _save_ctx: + ctx.save_for_backward(query, key, value) + ctx.attn_mask = attn_mask + ctx.dropout_p = dropout_p + ctx.is_causal = is_causal + ctx.scale = scale + ctx.enable_gqa = enable_gqa + + query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value)) + out = torch.nn.functional.scaled_dot_product_attention( + query=query, + key=key, + value=value, + attn_mask=attn_mask, + dropout_p=dropout_p, + is_causal=is_causal, + scale=scale, + enable_gqa=enable_gqa, + ) + out = out.permute(0, 2, 1, 3) + + return out + + +def _native_attention_backward_op( + ctx: torch.autograd.function.FunctionCtx, + grad_out: torch.Tensor, + *args, + **kwargs, +): + query, key, value = ctx.saved_tensors + + query.requires_grad_(True) + key.requires_grad_(True) + value.requires_grad_(True) + + query_t, key_t, value_t = (x.permute(0, 2, 1, 3) for x in (query, key, value)) + out = torch.nn.functional.scaled_dot_product_attention( + query=query_t, + key=key_t, + value=value_t, + attn_mask=ctx.attn_mask, + dropout_p=ctx.dropout_p, + is_causal=ctx.is_causal, + scale=ctx.scale, + enable_gqa=ctx.enable_gqa, + ) + out = out.permute(0, 2, 1, 3) + + grad_out_t = grad_out.permute(0, 2, 1, 3) + grad_query_t, grad_key_t, grad_value_t = torch.autograd.grad( + outputs=out, inputs=[query_t, key_t, value_t], grad_outputs=grad_out_t, retain_graph=False + ) + + grad_query = grad_query_t.permute(0, 2, 1, 3) + grad_key = grad_key_t.permute(0, 2, 1, 3) + grad_value = grad_value_t.permute(0, 2, 1, 3) + + return grad_query, grad_key, grad_value + + # https://github.com/pytorch/pytorch/blob/8904ba638726f8c9a5aff5977c4aa76c9d2edfa6/aten/src/ATen/native/native_functions.yaml#L14958 # forward declaration: # aten::_scaled_dot_product_cudnn_attention(Tensor query, Tensor key, Tensor value, Tensor? attn_bias, bool compute_log_sumexp, float dropout_p=0., bool is_causal=False, bool return_debug_mask=False, *, float? scale=None) -> (Tensor output, Tensor logsumexp, Tensor cum_seq_q, Tensor cum_seq_k, SymInt max_q, SymInt max_k, Tensor philox_seed, Tensor philox_offset, Tensor debug_attn_mask) @@ -1523,6 +1603,7 @@ def _native_flex_attention( @_AttentionBackendRegistry.register( AttentionBackendName.NATIVE, constraints=[_check_device, _check_shape], + supports_context_parallel=True, ) def _native_attention( query: torch.Tensor, @@ -1538,18 +1619,35 @@ def _native_attention( ) -> torch.Tensor: if return_lse: raise ValueError("Native attention backend does not support setting `return_lse=True`.") - query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value)) - out = torch.nn.functional.scaled_dot_product_attention( - query=query, - key=key, - value=value, - attn_mask=attn_mask, - dropout_p=dropout_p, - is_causal=is_causal, - scale=scale, - enable_gqa=enable_gqa, - ) - out = out.permute(0, 2, 1, 3) + if _parallel_config is None: + query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value)) + out = torch.nn.functional.scaled_dot_product_attention( + query=query, + key=key, + value=value, + attn_mask=attn_mask, + dropout_p=dropout_p, + is_causal=is_causal, + scale=scale, + enable_gqa=enable_gqa, + ) + out = out.permute(0, 2, 1, 3) + else: + out = _templated_context_parallel_attention( + query, + key, + value, + attn_mask, + dropout_p, + is_causal, + scale, + enable_gqa, + return_lse, + forward_op=_native_attention_forward_op, + backward_op=_native_attention_backward_op, + _parallel_config=_parallel_config, + ) + return out From 325a95051bb20787d3db9faba9b4f62b2a63c43a Mon Sep 17 00:00:00 2001 From: Lev Novitskiy <57654885+leffff@users.noreply.github.com> Date: Tue, 4 Nov 2025 03:38:07 +0300 Subject: [PATCH 116/123] Kandinsky 5.0 Docs fixes (#12582) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * add transformer pipeline first version * updates * fix 5sec generation * rewrite Kandinsky5T2VPipeline to diffusers style * add multiprompt support * remove prints in pipeline * add nabla attention * Wrap Transformer in Diffusers style * fix license * fix prompt type * add gradient checkpointing and peft support * add usage example * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: Álvaro Somoza * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: Álvaro Somoza * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: Álvaro Somoza * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: Álvaro Somoza * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: Álvaro Somoza * remove unused imports * add 10 second models support * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * remove no_grad and simplified prompt paddings * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * moved template to __init__ * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: YiYi Xu * moved sdps inside processor * remove oneline function * remove reset_dtype methods * Transformer: move all methods to forward * separated prompt encoding * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: YiYi Xu * refactoring * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: YiYi Xu * refactoring acording to https://github.com/huggingface/diffusers/commit/acabbc0033d4b4933fc651766a4aa026db2e6dc1 * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * Update src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py Co-authored-by: YiYi Xu * fixed * style +copies * Update src/diffusers/models/transformers/transformer_kandinsky.py Co-authored-by: Charles * more * Apply suggestions from code review * add lora loader doc * add compiled Nabla Attention * all needed changes for 10 sec models are added! * add docs * Apply style fixes * update docs * add kandinsky5 to toctree * add tests * fix tests * Apply style fixes * update tests * minor docs refactoring * refactor Kandinsky 5.0 Vide docs * Update docs/source/en/_toctree.yml --------- Co-authored-by: Álvaro Somoza Co-authored-by: YiYi Xu Co-authored-by: Charles Co-authored-by: Sayak Paul Co-authored-by: github-actions[bot] --- docs/source/en/_toctree.yml | 4 ++-- .../api/pipelines/{kandinsky5.md => kandinsky5_video.md} | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) rename docs/source/en/api/pipelines/{kandinsky5.md => kandinsky5_video.md} (90%) diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 8103c01643..251eb25899 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -529,8 +529,6 @@ title: Kandinsky 2.2 - local: api/pipelines/kandinsky3 title: Kandinsky 3 - - local: api/pipelines/kandinsky5 - title: Kandinsky 5 - local: api/pipelines/kolors title: Kolors - local: api/pipelines/latent_consistency_models @@ -656,6 +654,8 @@ title: Text2Video-Zero - local: api/pipelines/wan title: Wan + - local: api/pipelines/kandinsky5_video + title: Kandinsky 5.0 Video title: Video title: Pipelines - sections: diff --git a/docs/source/en/api/pipelines/kandinsky5.md b/docs/source/en/api/pipelines/kandinsky5_video.md similarity index 90% rename from docs/source/en/api/pipelines/kandinsky5.md rename to docs/source/en/api/pipelines/kandinsky5_video.md index a98a0826b7..533db23e1c 100644 --- a/docs/source/en/api/pipelines/kandinsky5.md +++ b/docs/source/en/api/pipelines/kandinsky5_video.md @@ -7,9 +7,9 @@ an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express o specific language governing permissions and limitations under the License. --> -# Kandinsky 5.0 +# Kandinsky 5.0 Video -Kandinsky 5.0 is created by the Kandinsky team: Alexey Letunovskiy, Maria Kovaleva, Ivan Kirillov, Lev Novitskiy, Denis Koposov, Dmitrii Mikhailov, Anna Averchenkova, Andrey Shutkin, Julia Agafonova, Olga Kim, Anastasiia Kargapoltseva, Nikita Kiselev, Anna Dmitrienko, Anastasia Maltseva, Kirill Chernyshev, Ilia Vasiliev, Viacheslav Vasilev, Vladimir Polovnikov, Yury Kolabushin, Alexander Belykh, Mikhail Mamaev, Anastasia Aliaskina, Tatiana Nikulina, Polina Gavrilova, Vladimir Arkhipkin, Vladimir Korviakov, Nikolai Gerasimenko, Denis Parkhomenko, Denis Dimitrov +Kandinsky 5.0 Video is created by the Kandinsky team: Alexey Letunovskiy, Maria Kovaleva, Ivan Kirillov, Lev Novitskiy, Denis Koposov, Dmitrii Mikhailov, Anna Averchenkova, Andrey Shutkin, Julia Agafonova, Olga Kim, Anastasiia Kargapoltseva, Nikita Kiselev, Anna Dmitrienko, Anastasia Maltseva, Kirill Chernyshev, Ilia Vasiliev, Viacheslav Vasilev, Vladimir Polovnikov, Yury Kolabushin, Alexander Belykh, Mikhail Mamaev, Anastasia Aliaskina, Tatiana Nikulina, Polina Gavrilova, Vladimir Arkhipkin, Vladimir Korviakov, Nikolai Gerasimenko, Denis Parkhomenko, Denis Dimitrov Kandinsky 5.0 is a family of diffusion models for Video & Image generation. Kandinsky 5.0 T2V Lite is a lightweight video generation model (2B parameters) that ranks #1 among open-source models in its class. It outperforms larger models and offers the best understanding of Russian concepts in the open-source ecosystem. @@ -92,7 +92,7 @@ pipe = pipe.to("cuda") pipe.transformer.set_attention_backend( "flex" -) # <--- Set attention backend to Flex +) # <--- Sett attention bakend to Flex pipe.transformer.compile( mode="max-autotune-no-cudagraphs", dynamic=True @@ -115,7 +115,7 @@ export_to_video(output, "output.mp4", fps=24, quality=9) ``` ### Diffusion Distilled model -**⚠️ Warning!** all nocfg and diffusion distilled models should be inferred without CFG (```guidance_scale=1.0```): +**⚠️ Warning!** all nocfg and diffusion distilled models should be infered wothout CFG (```guidance_scale=1.0```): ```python model_id = "ai-forever/Kandinsky-5.0-T2V-Lite-distilled16steps-5s-Diffusers" From ac5a1e28fc9cc233863bcfb2abb9eef6807f156f Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Tue, 4 Nov 2025 10:26:07 +0530 Subject: [PATCH 117/123] [docs] sort doc (#12586) sort doc --- docs/source/en/_toctree.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 251eb25899..5af95cba74 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -636,6 +636,8 @@ title: HunyuanVideo - local: api/pipelines/i2vgenxl title: I2VGen-XL + - local: api/pipelines/kandinsky5_video + title: Kandinsky 5.0 Video - local: api/pipelines/latte title: Latte - local: api/pipelines/ltx_video @@ -654,8 +656,6 @@ title: Text2Video-Zero - local: api/pipelines/wan title: Wan - - local: api/pipelines/kandinsky5_video - title: Kandinsky 5.0 Video title: Video title: Pipelines - sections: From dcfb18a2d340d8e1f0ff001b06d2931ffa8648da Mon Sep 17 00:00:00 2001 From: Linoy Tsaban <57615435+linoytsaban@users.noreply.github.com> Date: Tue, 4 Nov 2025 14:27:25 +0200 Subject: [PATCH 118/123] [LoRA] add support for more Qwen LoRAs (#12581) * fix bug when offload and cache_latents both enabled * fix --- src/diffusers/loaders/lora_conversion_utils.py | 4 ++++ src/diffusers/loaders/lora_pipeline.py | 3 ++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index 099dbfc1d2..2807416f97 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -2213,6 +2213,10 @@ def _convert_non_diffusers_qwen_lora_to_diffusers(state_dict): state_dict = {convert_key(k): v for k, v in state_dict.items()} + has_default = any("default." in k for k in state_dict) + if has_default: + state_dict = {k.replace("default.", ""): v for k, v in state_dict.items()} + converted_state_dict = {} all_keys = list(state_dict.keys()) down_key = ".lora_down.weight" diff --git a/src/diffusers/loaders/lora_pipeline.py b/src/diffusers/loaders/lora_pipeline.py index 2bb6c0ea02..25919a896a 100644 --- a/src/diffusers/loaders/lora_pipeline.py +++ b/src/diffusers/loaders/lora_pipeline.py @@ -4940,7 +4940,8 @@ class QwenImageLoraLoaderMixin(LoraBaseMixin): has_alphas_in_sd = any(k.endswith(".alpha") for k in state_dict) has_lora_unet = any(k.startswith("lora_unet_") for k in state_dict) has_diffusion_model = any(k.startswith("diffusion_model.") for k in state_dict) - if has_alphas_in_sd or has_lora_unet or has_diffusion_model: + has_default = any("default." in k for k in state_dict) + if has_alphas_in_sd or has_lora_unet or has_diffusion_model or has_default: state_dict = _convert_non_diffusers_qwen_lora_to_diffusers(state_dict) out = (state_dict, metadata) if return_lora_metadata else state_dict From 6198f8a12b0d456b2345a524671f0276e7f48fc1 Mon Sep 17 00:00:00 2001 From: Dhruv Nair Date: Thu, 6 Nov 2025 07:54:24 +0530 Subject: [PATCH 119/123] [Modular] Allow ModularPipeline to load from revisions (#12592) * update * update * update * update --------- Co-authored-by: YiYi Xu --- src/diffusers/utils/dynamic_modules_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/diffusers/utils/dynamic_modules_utils.py b/src/diffusers/utils/dynamic_modules_utils.py index b2ef5a29e0..1c0734cf35 100644 --- a/src/diffusers/utils/dynamic_modules_utils.py +++ b/src/diffusers/utils/dynamic_modules_utils.py @@ -358,6 +358,7 @@ def get_cached_module_file( proxies=proxies, local_files_only=local_files_only, local_dir=local_dir, + revision=revision, token=token, ) submodule = os.path.join("local", "--".join(pretrained_model_name_or_path.split("/"))) From 58f3771545c81fe6e9179149414d16d251d13d37 Mon Sep 17 00:00:00 2001 From: Joseph Turian Date: Thu, 6 Nov 2025 04:07:31 +0000 Subject: [PATCH 120/123] Add optional precision-preserving preprocessing for examples/unconditional_image_generation/train_unconditional.py (#12596) * Add optional precision-preserving preprocessing * Document decoder caveat for precision flag --------- Co-authored-by: Sayak Paul --- .../unconditional_image_generation/README.md | 2 + .../train_unconditional.py | 57 +++++++++++++++++-- 2 files changed, 53 insertions(+), 6 deletions(-) diff --git a/examples/unconditional_image_generation/README.md b/examples/unconditional_image_generation/README.md index 22f982509b..6f8276a632 100644 --- a/examples/unconditional_image_generation/README.md +++ b/examples/unconditional_image_generation/README.md @@ -104,6 +104,8 @@ To use your own dataset, there are 2 ways: - you can either provide your own folder as `--train_data_dir` - or you can upload your dataset to the hub (possibly as a private repo, if you prefer so), and simply pass the `--dataset_name` argument. +If your dataset contains 16 or 32-bit channels (for example, medical TIFFs), add the `--preserve_input_precision` flag so the preprocessing keeps the original precision while still training a 3-channel model. Precision still depends on the decoder: Pillow keeps 16-bit grayscale and float inputs, but many 16-bit RGB files are decoded as 8-bit RGB, and the flag cannot recover precision lost at load time. + Below, we explain both in more detail. #### Provide the dataset as a folder diff --git a/examples/unconditional_image_generation/train_unconditional.py b/examples/unconditional_image_generation/train_unconditional.py index 3ffeef1364..0cc96220b9 100644 --- a/examples/unconditional_image_generation/train_unconditional.py +++ b/examples/unconditional_image_generation/train_unconditional.py @@ -52,6 +52,24 @@ def _extract_into_tensor(arr, timesteps, broadcast_shape): return res.expand(broadcast_shape) +def _ensure_three_channels(tensor: torch.Tensor) -> torch.Tensor: + """ + Ensure the tensor has exactly three channels (C, H, W) by repeating or truncating channels when needed. + """ + if tensor.ndim == 2: + tensor = tensor.unsqueeze(0) + channels = tensor.shape[0] + if channels == 3: + return tensor + if channels == 1: + return tensor.repeat(3, 1, 1) + if channels == 2: + return torch.cat([tensor, tensor[:1]], dim=0) + if channels > 3: + return tensor[:3] + raise ValueError(f"Unsupported number of channels: {channels}") + + def parse_args(): parser = argparse.ArgumentParser(description="Simple example of a training script.") parser.add_argument( @@ -260,6 +278,11 @@ def parse_args(): parser.add_argument( "--enable_xformers_memory_efficient_attention", action="store_true", help="Whether or not to use xformers." ) + parser.add_argument( + "--preserve_input_precision", + action="store_true", + help="Preserve 16/32-bit image precision by avoiding 8-bit RGB conversion while still producing 3-channel tensors.", + ) args = parser.parse_args() env_local_rank = int(os.environ.get("LOCAL_RANK", -1)) @@ -453,19 +476,41 @@ def main(args): # https://huggingface.co/docs/datasets/v2.4.0/en/image_load#imagefolder # Preprocessing the datasets and DataLoaders creation. + spatial_augmentations = [ + transforms.Resize(args.resolution, interpolation=transforms.InterpolationMode.BILINEAR), + transforms.CenterCrop(args.resolution) if args.center_crop else transforms.RandomCrop(args.resolution), + transforms.RandomHorizontalFlip() if args.random_flip else transforms.Lambda(lambda x: x), + ] + augmentations = transforms.Compose( - [ - transforms.Resize(args.resolution, interpolation=transforms.InterpolationMode.BILINEAR), - transforms.CenterCrop(args.resolution) if args.center_crop else transforms.RandomCrop(args.resolution), - transforms.RandomHorizontalFlip() if args.random_flip else transforms.Lambda(lambda x: x), + spatial_augmentations + + [ transforms.ToTensor(), transforms.Normalize([0.5], [0.5]), ] ) + precision_augmentations = transforms.Compose( + [ + transforms.PILToTensor(), + transforms.Lambda(_ensure_three_channels), + transforms.ConvertImageDtype(torch.float32), + ] + + spatial_augmentations + + [transforms.Normalize([0.5], [0.5])] + ) + def transform_images(examples): - images = [augmentations(image.convert("RGB")) for image in examples["image"]] - return {"input": images} + processed = [] + for image in examples["image"]: + if not args.preserve_input_precision: + processed.append(augmentations(image.convert("RGB"))) + else: + precise_image = image + if precise_image.mode == "P": + precise_image = precise_image.convert("RGB") + processed.append(precision_augmentations(precise_image)) + return {"input": processed} logger.info(f"Dataset size: {len(dataset)}") From b3e9dfced7c9e8d00f646c710766b532383f04c6 Mon Sep 17 00:00:00 2001 From: Junsong Chen Date: Thu, 6 Nov 2025 13:08:47 +0800 Subject: [PATCH 121/123] [SANA-Video] Adding 5s pre-trained 480p SANA-Video inference (#12584) * 1. add `SanaVideoTransformer3DModel` in transformer_sana_video.py 2. add `SanaVideoPipeline` in pipeline_sana_video.py 3. add all code we need for import `SanaVideoPipeline` * add a sample about how to use sana-video; * code update; * update hf model path; * update code; * sana-video can run now; * 1. add aspect ratio in sana-video-pipeline; 2. add reshape function in sana-video-processor; 3. fix convert pth to safetensor bugs; * default to use `use_resolution_binning`; * make style; * remove unused code; * Update src/diffusers/models/transformers/transformer_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/models/transformers/transformer_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/models/transformers/transformer_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/pipelines/sana/pipeline_sana_video.py Co-authored-by: YiYi Xu * Update src/diffusers/models/transformers/transformer_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/models/transformers/transformer_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/models/transformers/transformer_sana_video.py * Update src/diffusers/pipelines/sana/pipeline_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/models/transformers/transformer_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/pipelines/sana/pipeline_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * support `dispatch_attention_fn` * 1. add sana-video markdown; 2. fix typos; * add two test case for sana-video (need check) * fix text-encoder in test-sana-video; * Update tests/pipelines/sana/test_sana_video.py * Update tests/pipelines/sana/test_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update tests/pipelines/sana/test_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update tests/pipelines/sana/test_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update tests/pipelines/sana/test_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update tests/pipelines/sana/test_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/pipelines/sana/pipeline_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/video_processor.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * make style make quality make fix-copies * toctree yaml update; * add sana-video-transformer3d markdown; * Apply style fixes --------- Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> Co-authored-by: YiYi Xu Co-authored-by: github-actions[bot] --- docs/source/en/_toctree.yml | 4 + .../en/api/models/sana_video_transformer3d.md | 36 + docs/source/en/api/pipelines/sana_sprint.md | 3 - docs/source/en/api/pipelines/sana_video.md | 102 ++ scripts/convert_sana_video_to_diffusers.py | 324 ++++++ src/diffusers/__init__.py | 4 + src/diffusers/models/__init__.py | 2 + src/diffusers/models/transformers/__init__.py | 1 + .../transformers/transformer_sana_video.py | 703 ++++++++++++ src/diffusers/pipelines/__init__.py | 9 +- src/diffusers/pipelines/sana/__init__.py | 2 + .../pipelines/sana/pipeline_output.py | 16 + src/diffusers/pipelines/sana/pipeline_sana.py | 2 +- .../pipelines/sana/pipeline_sana_sprint.py | 2 +- .../pipelines/sana/pipeline_sana_video.py | 1017 +++++++++++++++++ src/diffusers/utils/dummy_pt_objects.py | 15 + .../dummy_torch_and_transformers_objects.py | 15 + src/diffusers/video_processor.py | 65 +- .../test_models_transformer_sana_video.py | 97 ++ tests/pipelines/sana/test_sana_video.py | 225 ++++ 20 files changed, 2637 insertions(+), 7 deletions(-) create mode 100644 docs/source/en/api/models/sana_video_transformer3d.md create mode 100644 docs/source/en/api/pipelines/sana_video.md create mode 100644 scripts/convert_sana_video_to_diffusers.py create mode 100644 src/diffusers/models/transformers/transformer_sana_video.py create mode 100644 src/diffusers/pipelines/sana/pipeline_sana_video.py create mode 100644 tests/models/transformers/test_models_transformer_sana_video.py create mode 100644 tests/pipelines/sana/test_sana_video.py diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 5af95cba74..94dad286e4 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -373,6 +373,8 @@ title: QwenImageTransformer2DModel - local: api/models/sana_transformer2d title: SanaTransformer2DModel + - local: api/models/sana_video_transformer3d + title: SanaVideoTransformer3DModel - local: api/models/sd3_transformer2d title: SD3Transformer2DModel - local: api/models/skyreels_v2_transformer_3d @@ -563,6 +565,8 @@ title: Sana - local: api/pipelines/sana_sprint title: Sana Sprint + - local: api/pipelines/sana_video + title: Sana Video - local: api/pipelines/self_attention_guidance title: Self-Attention Guidance - local: api/pipelines/semantic_stable_diffusion diff --git a/docs/source/en/api/models/sana_video_transformer3d.md b/docs/source/en/api/models/sana_video_transformer3d.md new file mode 100644 index 0000000000..0cf1451a2d --- /dev/null +++ b/docs/source/en/api/models/sana_video_transformer3d.md @@ -0,0 +1,36 @@ + + +# SanaVideoTransformer3DModel + +A Diffusion Transformer model for 3D data (video) from [SANA-Video: Efficient Video Generation with Block Linear Diffusion Transformer](https://huggingface.co/papers/2509.24695) from NVIDIA and MIT HAN Lab, by Junsong Chen, Yuyang Zhao, Jincheng Yu, Ruihang Chu, Junyu Chen, Shuai Yang, Xianbang Wang, Yicheng Pan, Daquan Zhou, Huan Ling, Haozhe Liu, Hongwei Yi, Hao Zhang, Muyang Li, Yukang Chen, Han Cai, Sanja Fidler, Ping Luo, Song Han, Enze Xie. + +The abstract from the paper is: + +*We introduce SANA-Video, a small diffusion model that can efficiently generate videos up to 720x1280 resolution and minute-length duration. SANA-Video synthesizes high-resolution, high-quality and long videos with strong text-video alignment at a remarkably fast speed, deployable on RTX 5090 GPU. Two core designs ensure our efficient, effective and long video generation: (1) Linear DiT: We leverage linear attention as the core operation, which is more efficient than vanilla attention given the large number of tokens processed in video generation. (2) Constant-Memory KV cache for Block Linear Attention: we design block-wise autoregressive approach for long video generation by employing a constant-memory state, derived from the cumulative properties of linear attention. This KV cache provides the Linear DiT with global context at a fixed memory cost, eliminating the need for a traditional KV cache and enabling efficient, minute-long video generation. In addition, we explore effective data filters and model training strategies, narrowing the training cost to 12 days on 64 H100 GPUs, which is only 1% of the cost of MovieGen. Given its low cost, SANA-Video achieves competitive performance compared to modern state-of-the-art small diffusion models (e.g., Wan 2.1-1.3B and SkyReel-V2-1.3B) while being 16x faster in measured latency. Moreover, SANA-Video can be deployed on RTX 5090 GPUs with NVFP4 precision, accelerating the inference speed of generating a 5-second 720p video from 71s to 29s (2.4x speedup). In summary, SANA-Video enables low-cost, high-quality video generation.* + +The model can be loaded with the following code snippet. + +```python +from diffusers import SanaVideoTransformer3DModel +import torch + +transformer = SanaVideoTransformer3DModel.from_pretrained("Efficient-Large-Model/SANA-Video_2B_480p_diffusers", subfolder="transformer", torch_dtype=torch.bfloat16) +``` + +## SanaVideoTransformer3DModel + +[[autodoc]] SanaVideoTransformer3DModel + +## Transformer2DModelOutput + +[[autodoc]] models.modeling_outputs.Transformer2DModelOutput + diff --git a/docs/source/en/api/pipelines/sana_sprint.md b/docs/source/en/api/pipelines/sana_sprint.md index 357d7e406d..46cdc13302 100644 --- a/docs/source/en/api/pipelines/sana_sprint.md +++ b/docs/source/en/api/pipelines/sana_sprint.md @@ -24,9 +24,6 @@ The abstract from the paper is: *This paper presents SANA-Sprint, an efficient diffusion model for ultra-fast text-to-image (T2I) generation. SANA-Sprint is built on a pre-trained foundation model and augmented with hybrid distillation, dramatically reducing inference steps from 20 to 1-4. We introduce three key innovations: (1) We propose a training-free approach that transforms a pre-trained flow-matching model for continuous-time consistency distillation (sCM), eliminating costly training from scratch and achieving high training efficiency. Our hybrid distillation strategy combines sCM with latent adversarial distillation (LADD): sCM ensures alignment with the teacher model, while LADD enhances single-step generation fidelity. (2) SANA-Sprint is a unified step-adaptive model that achieves high-quality generation in 1-4 steps, eliminating step-specific training and improving efficiency. (3) We integrate ControlNet with SANA-Sprint for real-time interactive image generation, enabling instant visual feedback for user interaction. SANA-Sprint establishes a new Pareto frontier in speed-quality tradeoffs, achieving state-of-the-art performance with 7.59 FID and 0.74 GenEval in only 1 step — outperforming FLUX-schnell (7.94 FID / 0.71 GenEval) while being 10× faster (0.1s vs 1.1s on H100). It also achieves 0.1s (T2I) and 0.25s (ControlNet) latency for 1024×1024 images on H100, and 0.31s (T2I) on an RTX 4090, showcasing its exceptional efficiency and potential for AI-powered consumer applications (AIPC). Code and pre-trained models will be open-sourced.* -> [!TIP] -> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines. - This pipeline was contributed by [lawrence-cj](https://github.com/lawrence-cj), [shuchen Xue](https://github.com/scxue) and [Enze Xie](https://github.com/xieenze). The original codebase can be found [here](https://github.com/NVlabs/Sana). The original weights can be found under [hf.co/Efficient-Large-Model](https://huggingface.co/Efficient-Large-Model/). Available models: diff --git a/docs/source/en/api/pipelines/sana_video.md b/docs/source/en/api/pipelines/sana_video.md new file mode 100644 index 0000000000..85d77fb294 --- /dev/null +++ b/docs/source/en/api/pipelines/sana_video.md @@ -0,0 +1,102 @@ + + +# SanaVideoPipeline + +
+ LoRA + MPS +
+ +[SANA-Video: Efficient Video Generation with Block Linear Diffusion Transformer](https://huggingface.co/papers/2509.24695) from NVIDIA and MIT HAN Lab, by Junsong Chen, Yuyang Zhao, Jincheng Yu, Ruihang Chu, Junyu Chen, Shuai Yang, Xianbang Wang, Yicheng Pan, Daquan Zhou, Huan Ling, Haozhe Liu, Hongwei Yi, Hao Zhang, Muyang Li, Yukang Chen, Han Cai, Sanja Fidler, Ping Luo, Song Han, Enze Xie. + +The abstract from the paper is: + +*We introduce SANA-Video, a small diffusion model that can efficiently generate videos up to 720x1280 resolution and minute-length duration. SANA-Video synthesizes high-resolution, high-quality and long videos with strong text-video alignment at a remarkably fast speed, deployable on RTX 5090 GPU. Two core designs ensure our efficient, effective and long video generation: (1) Linear DiT: We leverage linear attention as the core operation, which is more efficient than vanilla attention given the large number of tokens processed in video generation. (2) Constant-Memory KV cache for Block Linear Attention: we design block-wise autoregressive approach for long video generation by employing a constant-memory state, derived from the cumulative properties of linear attention. This KV cache provides the Linear DiT with global context at a fixed memory cost, eliminating the need for a traditional KV cache and enabling efficient, minute-long video generation. In addition, we explore effective data filters and model training strategies, narrowing the training cost to 12 days on 64 H100 GPUs, which is only 1% of the cost of MovieGen. Given its low cost, SANA-Video achieves competitive performance compared to modern state-of-the-art small diffusion models (e.g., Wan 2.1-1.3B and SkyReel-V2-1.3B) while being 16x faster in measured latency. Moreover, SANA-Video can be deployed on RTX 5090 GPUs with NVFP4 precision, accelerating the inference speed of generating a 5-second 720p video from 71s to 29s (2.4x speedup). In summary, SANA-Video enables low-cost, high-quality video generation. [this https URL](https://github.com/NVlabs/SANA).* + +This pipeline was contributed by SANA Team. The original codebase can be found [here](https://github.com/NVlabs/Sana). The original weights can be found under [hf.co/Efficient-Large-Model](https://hf.co/collections/Efficient-Large-Model/sana-video). + +Available models: + +| Model | Recommended dtype | +|:-----:|:-----------------:| +| [`Efficient-Large-Model/SANA-Video_2B_480p_diffusers`](https://huggingface.co/Efficient-Large-Model/ANA-Video_2B_480p_diffusers) | `torch.bfloat16` | + +Refer to [this](https://huggingface.co/collections/Efficient-Large-Model/sana-video) collection for more information. + +Note: The recommended dtype mentioned is for the transformer weights. The text encoder and VAE weights must stay in `torch.bfloat16` or `torch.float32` for the model to work correctly. Please refer to the inference example below to see how to load the model with the recommended dtype. + +## Quantization + +Quantization helps reduce the memory requirements of very large models by storing model weights in a lower precision data type. However, quantization may have varying impact on video quality depending on the video model. + +Refer to the [Quantization](../../quantization/overview) overview to learn more about supported quantization backends and selecting a quantization backend that supports your use case. The example below demonstrates how to load a quantized [`SanaVideoPipeline`] for inference with bitsandbytes. + +```py +import torch +from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig, SanaVideoTransformer3DModel, SanaVideoPipeline +from transformers import BitsAndBytesConfig as BitsAndBytesConfig, AutoModel + +quant_config = BitsAndBytesConfig(load_in_8bit=True) +text_encoder_8bit = AutoModel.from_pretrained( + "Efficient-Large-Model/SANA-Video_2B_480p_diffusers", + subfolder="text_encoder", + quantization_config=quant_config, + torch_dtype=torch.float16, +) + +quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True) +transformer_8bit = SanaVideoTransformer3DModel.from_pretrained( + "Efficient-Large-Model/SANA-Video_2B_480p_diffusers", + subfolder="transformer", + quantization_config=quant_config, + torch_dtype=torch.float16, +) + +pipeline = SanaVideoPipeline.from_pretrained( + "Efficient-Large-Model/SANA-Video_2B_480p_diffusers", + text_encoder=text_encoder_8bit, + transformer=transformer_8bit, + torch_dtype=torch.float16, + device_map="balanced", +) + +model_score = 30 +prompt = "Evening, backlight, side lighting, soft light, high contrast, mid-shot, centered composition, clean solo shot, warm color. A young Caucasian man stands in a forest, golden light glimmers on his hair as sunlight filters through the leaves. He wears a light shirt, wind gently blowing his hair and collar, light dances across his face with his movements. The background is blurred, with dappled light and soft tree shadows in the distance. The camera focuses on his lifted gaze, clear and emotional." +negative_prompt = "A chaotic sequence with misshapen, deformed limbs in heavy motion blur, sudden disappearance, jump cuts, jerky movements, rapid shot changes, frames out of sync, inconsistent character shapes, temporal artifacts, jitter, and ghosting effects, creating a disorienting visual experience." +motion_prompt = f" motion score: {model_score}." +prompt = prompt + motion_prompt + +output = pipeline( + prompt=prompt, + negative_prompt=negative_prompt, + height=480, + width=832, + num_frames=81, + guidance_scale=6.0, + num_inference_steps=50 +).frames[0] +export_to_video(output, "sana-video-output.mp4", fps=16) +``` + +## SanaVideoPipeline + +[[autodoc]] SanaVideoPipeline + - all + - __call__ + + +## SanaVideoPipelineOutput + +[[autodoc]] pipelines.sana.pipeline_sana_video.SanaVideoPipelineOutput diff --git a/scripts/convert_sana_video_to_diffusers.py b/scripts/convert_sana_video_to_diffusers.py new file mode 100644 index 0000000000..fbb7c1d9e7 --- /dev/null +++ b/scripts/convert_sana_video_to_diffusers.py @@ -0,0 +1,324 @@ +#!/usr/bin/env python +from __future__ import annotations + +import argparse +import os +from contextlib import nullcontext + +import torch +from accelerate import init_empty_weights +from huggingface_hub import hf_hub_download, snapshot_download +from termcolor import colored +from transformers import AutoModelForCausalLM, AutoTokenizer + +from diffusers import ( + AutoencoderKLWan, + DPMSolverMultistepScheduler, + FlowMatchEulerDiscreteScheduler, + SanaVideoPipeline, + SanaVideoTransformer3DModel, + UniPCMultistepScheduler, +) +from diffusers.utils.import_utils import is_accelerate_available + + +CTX = init_empty_weights if is_accelerate_available else nullcontext + +ckpt_ids = ["Efficient-Large-Model/SANA-Video_2B_480p/checkpoints/SANA_Video_2B_480p.pth"] +# https://github.com/NVlabs/Sana/blob/main/inference_video_scripts/inference_sana_video.py + + +def main(args): + cache_dir_path = os.path.expanduser("~/.cache/huggingface/hub") + + if args.orig_ckpt_path is None or args.orig_ckpt_path in ckpt_ids: + ckpt_id = args.orig_ckpt_path or ckpt_ids[0] + snapshot_download( + repo_id=f"{'/'.join(ckpt_id.split('/')[:2])}", + cache_dir=cache_dir_path, + repo_type="model", + ) + file_path = hf_hub_download( + repo_id=f"{'/'.join(ckpt_id.split('/')[:2])}", + filename=f"{'/'.join(ckpt_id.split('/')[2:])}", + cache_dir=cache_dir_path, + repo_type="model", + ) + else: + file_path = args.orig_ckpt_path + + print(colored(f"Loading checkpoint from {file_path}", "green", attrs=["bold"])) + all_state_dict = torch.load(file_path, weights_only=True) + state_dict = all_state_dict.pop("state_dict") + converted_state_dict = {} + + # Patch embeddings. + converted_state_dict["patch_embedding.weight"] = state_dict.pop("x_embedder.proj.weight") + converted_state_dict["patch_embedding.bias"] = state_dict.pop("x_embedder.proj.bias") + + # Caption projection. + converted_state_dict["caption_projection.linear_1.weight"] = state_dict.pop("y_embedder.y_proj.fc1.weight") + converted_state_dict["caption_projection.linear_1.bias"] = state_dict.pop("y_embedder.y_proj.fc1.bias") + converted_state_dict["caption_projection.linear_2.weight"] = state_dict.pop("y_embedder.y_proj.fc2.weight") + converted_state_dict["caption_projection.linear_2.bias"] = state_dict.pop("y_embedder.y_proj.fc2.bias") + + converted_state_dict["time_embed.emb.timestep_embedder.linear_1.weight"] = state_dict.pop( + "t_embedder.mlp.0.weight" + ) + converted_state_dict["time_embed.emb.timestep_embedder.linear_1.bias"] = state_dict.pop("t_embedder.mlp.0.bias") + converted_state_dict["time_embed.emb.timestep_embedder.linear_2.weight"] = state_dict.pop( + "t_embedder.mlp.2.weight" + ) + converted_state_dict["time_embed.emb.timestep_embedder.linear_2.bias"] = state_dict.pop("t_embedder.mlp.2.bias") + + # Shared norm. + converted_state_dict["time_embed.linear.weight"] = state_dict.pop("t_block.1.weight") + converted_state_dict["time_embed.linear.bias"] = state_dict.pop("t_block.1.bias") + + # y norm + converted_state_dict["caption_norm.weight"] = state_dict.pop("attention_y_norm.weight") + + # scheduler + flow_shift = 8.0 + + # model config + layer_num = 20 + # Positional embedding interpolation scale. + qk_norm = True + + # sample size + if args.video_size == 480: + sample_size = 30 # Wan-VAE: 8xp2 downsample factor + patch_size = (1, 2, 2) + elif args.video_size == 720: + sample_size = 22 # Wan-VAE: 32xp1 downsample factor + patch_size = (1, 1, 1) + else: + raise ValueError(f"Video size {args.video_size} is not supported.") + + for depth in range(layer_num): + # Transformer blocks. + converted_state_dict[f"transformer_blocks.{depth}.scale_shift_table"] = state_dict.pop( + f"blocks.{depth}.scale_shift_table" + ) + + # Linear Attention is all you need 🤘 + # Self attention. + q, k, v = torch.chunk(state_dict.pop(f"blocks.{depth}.attn.qkv.weight"), 3, dim=0) + converted_state_dict[f"transformer_blocks.{depth}.attn1.to_q.weight"] = q + converted_state_dict[f"transformer_blocks.{depth}.attn1.to_k.weight"] = k + converted_state_dict[f"transformer_blocks.{depth}.attn1.to_v.weight"] = v + if qk_norm is not None: + # Add Q/K normalization for self-attention (attn1) - needed for Sana-Sprint and Sana-1.5 + converted_state_dict[f"transformer_blocks.{depth}.attn1.norm_q.weight"] = state_dict.pop( + f"blocks.{depth}.attn.q_norm.weight" + ) + converted_state_dict[f"transformer_blocks.{depth}.attn1.norm_k.weight"] = state_dict.pop( + f"blocks.{depth}.attn.k_norm.weight" + ) + # Projection. + converted_state_dict[f"transformer_blocks.{depth}.attn1.to_out.0.weight"] = state_dict.pop( + f"blocks.{depth}.attn.proj.weight" + ) + converted_state_dict[f"transformer_blocks.{depth}.attn1.to_out.0.bias"] = state_dict.pop( + f"blocks.{depth}.attn.proj.bias" + ) + + # Feed-forward. + converted_state_dict[f"transformer_blocks.{depth}.ff.conv_inverted.weight"] = state_dict.pop( + f"blocks.{depth}.mlp.inverted_conv.conv.weight" + ) + converted_state_dict[f"transformer_blocks.{depth}.ff.conv_inverted.bias"] = state_dict.pop( + f"blocks.{depth}.mlp.inverted_conv.conv.bias" + ) + converted_state_dict[f"transformer_blocks.{depth}.ff.conv_depth.weight"] = state_dict.pop( + f"blocks.{depth}.mlp.depth_conv.conv.weight" + ) + converted_state_dict[f"transformer_blocks.{depth}.ff.conv_depth.bias"] = state_dict.pop( + f"blocks.{depth}.mlp.depth_conv.conv.bias" + ) + converted_state_dict[f"transformer_blocks.{depth}.ff.conv_point.weight"] = state_dict.pop( + f"blocks.{depth}.mlp.point_conv.conv.weight" + ) + converted_state_dict[f"transformer_blocks.{depth}.ff.conv_temp.weight"] = state_dict.pop( + f"blocks.{depth}.mlp.t_conv.weight" + ) + + # Cross-attention. + q = state_dict.pop(f"blocks.{depth}.cross_attn.q_linear.weight") + q_bias = state_dict.pop(f"blocks.{depth}.cross_attn.q_linear.bias") + k, v = torch.chunk(state_dict.pop(f"blocks.{depth}.cross_attn.kv_linear.weight"), 2, dim=0) + k_bias, v_bias = torch.chunk(state_dict.pop(f"blocks.{depth}.cross_attn.kv_linear.bias"), 2, dim=0) + + converted_state_dict[f"transformer_blocks.{depth}.attn2.to_q.weight"] = q + converted_state_dict[f"transformer_blocks.{depth}.attn2.to_q.bias"] = q_bias + converted_state_dict[f"transformer_blocks.{depth}.attn2.to_k.weight"] = k + converted_state_dict[f"transformer_blocks.{depth}.attn2.to_k.bias"] = k_bias + converted_state_dict[f"transformer_blocks.{depth}.attn2.to_v.weight"] = v + converted_state_dict[f"transformer_blocks.{depth}.attn2.to_v.bias"] = v_bias + if qk_norm is not None: + # Add Q/K normalization for cross-attention (attn2) - needed for Sana-Sprint and Sana-1.5 + converted_state_dict[f"transformer_blocks.{depth}.attn2.norm_q.weight"] = state_dict.pop( + f"blocks.{depth}.cross_attn.q_norm.weight" + ) + converted_state_dict[f"transformer_blocks.{depth}.attn2.norm_k.weight"] = state_dict.pop( + f"blocks.{depth}.cross_attn.k_norm.weight" + ) + + converted_state_dict[f"transformer_blocks.{depth}.attn2.to_out.0.weight"] = state_dict.pop( + f"blocks.{depth}.cross_attn.proj.weight" + ) + converted_state_dict[f"transformer_blocks.{depth}.attn2.to_out.0.bias"] = state_dict.pop( + f"blocks.{depth}.cross_attn.proj.bias" + ) + + # Final block. + converted_state_dict["proj_out.weight"] = state_dict.pop("final_layer.linear.weight") + converted_state_dict["proj_out.bias"] = state_dict.pop("final_layer.linear.bias") + converted_state_dict["scale_shift_table"] = state_dict.pop("final_layer.scale_shift_table") + + # Transformer + with CTX(): + transformer_kwargs = { + "in_channels": 16, + "out_channels": 16, + "num_attention_heads": 20, + "attention_head_dim": 112, + "num_layers": 20, + "num_cross_attention_heads": 20, + "cross_attention_head_dim": 112, + "cross_attention_dim": 2240, + "caption_channels": 2304, + "mlp_ratio": 3.0, + "attention_bias": False, + "sample_size": sample_size, + "patch_size": patch_size, + "norm_elementwise_affine": False, + "norm_eps": 1e-6, + "qk_norm": "rms_norm_across_heads", + "rope_max_seq_len": 1024, + } + + transformer = SanaVideoTransformer3DModel(**transformer_kwargs) + + transformer.load_state_dict(converted_state_dict, strict=True, assign=True) + + try: + state_dict.pop("y_embedder.y_embedding") + state_dict.pop("pos_embed") + state_dict.pop("logvar_linear.weight") + state_dict.pop("logvar_linear.bias") + except KeyError: + print("y_embedder.y_embedding or pos_embed not found in the state_dict") + + assert len(state_dict) == 0, f"State dict is not empty, {state_dict.keys()}" + + num_model_params = sum(p.numel() for p in transformer.parameters()) + print(f"Total number of transformer parameters: {num_model_params}") + + transformer = transformer.to(weight_dtype) + + if not args.save_full_pipeline: + print( + colored( + f"Only saving transformer model of {args.model_type}. " + f"Set --save_full_pipeline to save the whole Pipeline", + "green", + attrs=["bold"], + ) + ) + transformer.save_pretrained( + os.path.join(args.dump_path, "transformer"), safe_serialization=True, max_shard_size="5GB" + ) + else: + print(colored(f"Saving the whole Pipeline containing {args.model_type}", "green", attrs=["bold"])) + # VAE + vae = AutoencoderKLWan.from_pretrained( + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", subfolder="vae", torch_dtype=torch.float32 + ) + + # Text Encoder + text_encoder_model_path = "Efficient-Large-Model/gemma-2-2b-it" + tokenizer = AutoTokenizer.from_pretrained(text_encoder_model_path) + tokenizer.padding_side = "right" + text_encoder = AutoModelForCausalLM.from_pretrained( + text_encoder_model_path, torch_dtype=torch.bfloat16 + ).get_decoder() + + # Choose the appropriate pipeline and scheduler based on model type + # Original Sana scheduler + if args.scheduler_type == "flow-dpm_solver": + scheduler = DPMSolverMultistepScheduler( + flow_shift=flow_shift, + use_flow_sigmas=True, + prediction_type="flow_prediction", + ) + elif args.scheduler_type == "flow-euler": + scheduler = FlowMatchEulerDiscreteScheduler(shift=flow_shift) + elif args.scheduler_type == "uni-pc": + scheduler = UniPCMultistepScheduler( + prediction_type="flow_prediction", + use_flow_sigmas=True, + num_train_timesteps=1000, + flow_shift=flow_shift, + ) + else: + raise ValueError(f"Scheduler type {args.scheduler_type} is not supported") + + pipe = SanaVideoPipeline( + tokenizer=tokenizer, + text_encoder=text_encoder, + transformer=transformer, + vae=vae, + scheduler=scheduler, + ) + + pipe.save_pretrained(args.dump_path, safe_serialization=True, max_shard_size="5GB") + + +DTYPE_MAPPING = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, +} + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--orig_ckpt_path", default=None, type=str, required=False, help="Path to the checkpoint to convert." + ) + parser.add_argument( + "--video_size", + default=480, + type=int, + choices=[480, 720], + required=False, + help="Video size of pretrained model, 480 or 720.", + ) + parser.add_argument( + "--model_type", + default="SanaVideo", + type=str, + choices=[ + "SanaVideo", + ], + ) + parser.add_argument( + "--scheduler_type", + default="flow-dpm_solver", + type=str, + choices=["flow-dpm_solver", "flow-euler", "uni-pc"], + help="Scheduler type to use.", + ) + parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output pipeline.") + parser.add_argument("--save_full_pipeline", action="store_true", help="save all the pipeline elements in one.") + parser.add_argument("--dtype", default="fp32", type=str, choices=["fp32", "fp16", "bf16"], help="Weight dtype.") + + args = parser.parse_args() + + device = "cuda" if torch.cuda.is_available() else "cpu" + weight_dtype = DTYPE_MAPPING[args.dtype] + + main(args) diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 94104667b5..572aad4bd3 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -246,6 +246,7 @@ else: "QwenImageTransformer2DModel", "SanaControlNetModel", "SanaTransformer2DModel", + "SanaVideoTransformer3DModel", "SD3ControlNetModel", "SD3MultiControlNetModel", "SD3Transformer2DModel", @@ -544,6 +545,7 @@ else: "SanaPipeline", "SanaSprintImg2ImgPipeline", "SanaSprintPipeline", + "SanaVideoPipeline", "SemanticStableDiffusionPipeline", "ShapEImg2ImgPipeline", "ShapEPipeline", @@ -951,6 +953,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: QwenImageTransformer2DModel, SanaControlNetModel, SanaTransformer2DModel, + SanaVideoTransformer3DModel, SD3ControlNetModel, SD3MultiControlNetModel, SD3Transformer2DModel, @@ -1219,6 +1222,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: SanaPipeline, SanaSprintImg2ImgPipeline, SanaSprintPipeline, + SanaVideoPipeline, SemanticStableDiffusionPipeline, ShapEImg2ImgPipeline, ShapEPipeline, diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index e3b2974641..202e77fd19 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -102,6 +102,7 @@ if is_torch_available(): _import_structure["transformers.transformer_omnigen"] = ["OmniGenTransformer2DModel"] _import_structure["transformers.transformer_prx"] = ["PRXTransformer2DModel"] _import_structure["transformers.transformer_qwenimage"] = ["QwenImageTransformer2DModel"] + _import_structure["transformers.transformer_sana_video"] = ["SanaVideoTransformer3DModel"] _import_structure["transformers.transformer_sd3"] = ["SD3Transformer2DModel"] _import_structure["transformers.transformer_skyreels_v2"] = ["SkyReelsV2Transformer3DModel"] _import_structure["transformers.transformer_temporal"] = ["TransformerTemporalModel"] @@ -204,6 +205,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: PRXTransformer2DModel, QwenImageTransformer2DModel, SanaTransformer2DModel, + SanaVideoTransformer3DModel, SD3Transformer2DModel, SkyReelsV2Transformer3DModel, StableAudioDiTModel, diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index 2fe1159eec..15408a4b15 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -36,6 +36,7 @@ if is_torch_available(): from .transformer_omnigen import OmniGenTransformer2DModel from .transformer_prx import PRXTransformer2DModel from .transformer_qwenimage import QwenImageTransformer2DModel + from .transformer_sana_video import SanaVideoTransformer3DModel from .transformer_sd3 import SD3Transformer2DModel from .transformer_skyreels_v2 import SkyReelsV2Transformer3DModel from .transformer_temporal import TransformerTemporalModel diff --git a/src/diffusers/models/transformers/transformer_sana_video.py b/src/diffusers/models/transformers/transformer_sana_video.py new file mode 100644 index 0000000000..aaf96175c0 --- /dev/null +++ b/src/diffusers/models/transformers/transformer_sana_video.py @@ -0,0 +1,703 @@ +# Copyright 2025 The HuggingFace Team and SANA-Video Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from typing import Any, Dict, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +from torch import nn + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin, PeftAdapterMixin +from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers +from ..attention import AttentionMixin +from ..attention_dispatch import dispatch_attention_fn +from ..attention_processor import Attention +from ..embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps, get_1d_rotary_pos_embed +from ..modeling_outputs import Transformer2DModelOutput +from ..modeling_utils import ModelMixin +from ..normalization import AdaLayerNormSingle, RMSNorm + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class GLUMBTempConv(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + expand_ratio: float = 4, + norm_type: Optional[str] = None, + residual_connection: bool = True, + ) -> None: + super().__init__() + + hidden_channels = int(expand_ratio * in_channels) + self.norm_type = norm_type + self.residual_connection = residual_connection + + self.nonlinearity = nn.SiLU() + self.conv_inverted = nn.Conv2d(in_channels, hidden_channels * 2, 1, 1, 0) + self.conv_depth = nn.Conv2d(hidden_channels * 2, hidden_channels * 2, 3, 1, 1, groups=hidden_channels * 2) + self.conv_point = nn.Conv2d(hidden_channels, out_channels, 1, 1, 0, bias=False) + + self.norm = None + if norm_type == "rms_norm": + self.norm = RMSNorm(out_channels, eps=1e-5, elementwise_affine=True, bias=True) + + self.conv_temp = nn.Conv2d( + out_channels, out_channels, kernel_size=(3, 1), stride=1, padding=(1, 0), bias=False + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.residual_connection: + residual = hidden_states + batch_size, num_frames, height, width, num_channels = hidden_states.shape + hidden_states = hidden_states.view(batch_size * num_frames, height, width, num_channels).permute(0, 3, 1, 2) + + hidden_states = self.conv_inverted(hidden_states) + hidden_states = self.nonlinearity(hidden_states) + + hidden_states = self.conv_depth(hidden_states) + hidden_states, gate = torch.chunk(hidden_states, 2, dim=1) + hidden_states = hidden_states * self.nonlinearity(gate) + + hidden_states = self.conv_point(hidden_states) + + # Temporal aggregation + hidden_states_temporal = hidden_states.view(batch_size, num_frames, num_channels, height * width).permute( + 0, 2, 1, 3 + ) + hidden_states = hidden_states_temporal + self.conv_temp(hidden_states_temporal) + hidden_states = hidden_states.permute(0, 2, 3, 1).view(batch_size, num_frames, height, width, num_channels) + + if self.norm_type == "rms_norm": + # move channel to the last dimension so we apply RMSnorm across channel dimension + hidden_states = self.norm(hidden_states.movedim(1, -1)).movedim(-1, 1) + + if self.residual_connection: + hidden_states = hidden_states + residual + + return hidden_states + + +class SanaLinearAttnProcessor3_0: + r""" + Processor for implementing scaled dot-product linear attention. + """ + + def __call__( + self, + attn: Attention, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + rotary_emb: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + original_dtype = hidden_states.dtype + + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + + query = attn.to_q(hidden_states) + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) + + if attn.norm_q is not None: + query = attn.norm_q(query) + if attn.norm_k is not None: + key = attn.norm_k(key) + + query = query.unflatten(2, (attn.heads, -1)) + key = key.unflatten(2, (attn.heads, -1)) + value = value.unflatten(2, (attn.heads, -1)) + # B,N,H,C + + query = F.relu(query) + key = F.relu(key) + + if rotary_emb is not None: + + def apply_rotary_emb( + hidden_states: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + ): + x1, x2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1) + cos = freqs_cos[..., 0::2] + sin = freqs_sin[..., 1::2] + out = torch.empty_like(hidden_states) + out[..., 0::2] = x1 * cos - x2 * sin + out[..., 1::2] = x1 * sin + x2 * cos + return out.type_as(hidden_states) + + query_rotate = apply_rotary_emb(query, *rotary_emb) + key_rotate = apply_rotary_emb(key, *rotary_emb) + + # B,H,C,N + query = query.permute(0, 2, 3, 1) + key = key.permute(0, 2, 3, 1) + query_rotate = query_rotate.permute(0, 2, 3, 1) + key_rotate = key_rotate.permute(0, 2, 3, 1) + value = value.permute(0, 2, 3, 1) + + query_rotate, key_rotate, value = query_rotate.float(), key_rotate.float(), value.float() + + z = 1 / (key.sum(dim=-1, keepdim=True).transpose(-2, -1) @ query + 1e-15) + + scores = torch.matmul(value, key_rotate.transpose(-1, -2)) + hidden_states = torch.matmul(scores, query_rotate) + + hidden_states = hidden_states * z + # B,H,C,N + hidden_states = hidden_states.flatten(1, 2).transpose(1, 2) + hidden_states = hidden_states.to(original_dtype) + + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + + return hidden_states + + +# Copied from diffusers.models.transformers.transformer_wan.WanRotaryPosEmbed +class WanRotaryPosEmbed(nn.Module): + def __init__( + self, + attention_head_dim: int, + patch_size: Tuple[int, int, int], + max_seq_len: int, + theta: float = 10000.0, + ): + super().__init__() + + self.attention_head_dim = attention_head_dim + self.patch_size = patch_size + self.max_seq_len = max_seq_len + + h_dim = w_dim = 2 * (attention_head_dim // 6) + t_dim = attention_head_dim - h_dim - w_dim + freqs_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64 + + freqs_cos = [] + freqs_sin = [] + + for dim in [t_dim, h_dim, w_dim]: + freq_cos, freq_sin = get_1d_rotary_pos_embed( + dim, + max_seq_len, + theta, + use_real=True, + repeat_interleave_real=True, + freqs_dtype=freqs_dtype, + ) + freqs_cos.append(freq_cos) + freqs_sin.append(freq_sin) + + self.register_buffer("freqs_cos", torch.cat(freqs_cos, dim=1), persistent=False) + self.register_buffer("freqs_sin", torch.cat(freqs_sin, dim=1), persistent=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p_t, p_h, p_w = self.patch_size + ppf, pph, ppw = num_frames // p_t, height // p_h, width // p_w + + split_sizes = [ + self.attention_head_dim - 2 * (self.attention_head_dim // 3), + self.attention_head_dim // 3, + self.attention_head_dim // 3, + ] + + freqs_cos = self.freqs_cos.split(split_sizes, dim=1) + freqs_sin = self.freqs_sin.split(split_sizes, dim=1) + + freqs_cos_f = freqs_cos[0][:ppf].view(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1) + freqs_cos_h = freqs_cos[1][:pph].view(1, pph, 1, -1).expand(ppf, pph, ppw, -1) + freqs_cos_w = freqs_cos[2][:ppw].view(1, 1, ppw, -1).expand(ppf, pph, ppw, -1) + + freqs_sin_f = freqs_sin[0][:ppf].view(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1) + freqs_sin_h = freqs_sin[1][:pph].view(1, pph, 1, -1).expand(ppf, pph, ppw, -1) + freqs_sin_w = freqs_sin[2][:ppw].view(1, 1, ppw, -1).expand(ppf, pph, ppw, -1) + + freqs_cos = torch.cat([freqs_cos_f, freqs_cos_h, freqs_cos_w], dim=-1).reshape(1, ppf * pph * ppw, 1, -1) + freqs_sin = torch.cat([freqs_sin_f, freqs_sin_h, freqs_sin_w], dim=-1).reshape(1, ppf * pph * ppw, 1, -1) + + return freqs_cos, freqs_sin + + +# Copied from diffusers.models.transformers.sana_transformer.SanaModulatedNorm +class SanaModulatedNorm(nn.Module): + def __init__(self, dim: int, elementwise_affine: bool = False, eps: float = 1e-6): + super().__init__() + self.norm = nn.LayerNorm(dim, elementwise_affine=elementwise_affine, eps=eps) + + def forward( + self, hidden_states: torch.Tensor, temb: torch.Tensor, scale_shift_table: torch.Tensor + ) -> torch.Tensor: + hidden_states = self.norm(hidden_states) + shift, scale = (scale_shift_table[None] + temb[:, None].to(scale_shift_table.device)).chunk(2, dim=1) + hidden_states = hidden_states * (1 + scale) + shift + return hidden_states + + +class SanaCombinedTimestepGuidanceEmbeddings(nn.Module): + def __init__(self, embedding_dim): + super().__init__() + self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + self.guidance_condition_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.guidance_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=True) + + def forward(self, timestep: torch.Tensor, guidance: torch.Tensor = None, hidden_dtype: torch.dtype = None): + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, D) + + guidance_proj = self.guidance_condition_proj(guidance) + guidance_emb = self.guidance_embedder(guidance_proj.to(dtype=hidden_dtype)) + conditioning = timesteps_emb + guidance_emb + + return self.linear(self.silu(conditioning)), conditioning + + +class SanaAttnProcessor2_0: + r""" + Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). + """ + + _attention_backend = None + _parallel_config = None + + def __init__(self): + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError("SanaAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") + + def __call__( + self, + attn: Attention, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + batch_size, sequence_length, _ = ( + hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape + ) + + if attention_mask is not None: + attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) + # scaled_dot_product_attention expects attention_mask shape to be + # (batch, heads, source_length, target_length) + attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) + + query = attn.to_q(hidden_states) + + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) + + if attn.norm_q is not None: + query = attn.norm_q(query) + if attn.norm_k is not None: + key = attn.norm_k(key) + + inner_dim = key.shape[-1] + head_dim = inner_dim // attn.heads + + query = query.view(batch_size, -1, attn.heads, head_dim) + key = key.view(batch_size, -1, attn.heads, head_dim) + value = value.view(batch_size, -1, attn.heads, head_dim) + + # the output of sdp = (batch, num_heads, seq_len, head_dim) + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.type_as(query) + + # linear proj + hidden_states = attn.to_out[0](hidden_states) + # dropout + hidden_states = attn.to_out[1](hidden_states) + + hidden_states = hidden_states / attn.rescale_output_factor + + return hidden_states + + +class SanaVideoTransformerBlock(nn.Module): + r""" + Transformer block introduced in [Sana-Video](https://huggingface.co/papers/2509.24695). + """ + + def __init__( + self, + dim: int = 2240, + num_attention_heads: int = 20, + attention_head_dim: int = 112, + dropout: float = 0.0, + num_cross_attention_heads: Optional[int] = 20, + cross_attention_head_dim: Optional[int] = 112, + cross_attention_dim: Optional[int] = 2240, + attention_bias: bool = True, + norm_elementwise_affine: bool = False, + norm_eps: float = 1e-6, + attention_out_bias: bool = True, + mlp_ratio: float = 3.0, + qk_norm: Optional[str] = "rms_norm_across_heads", + rope_max_seq_len: int = 1024, + ) -> None: + super().__init__() + + # 1. Self Attention + self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=norm_eps) + self.attn1 = Attention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + kv_heads=num_attention_heads if qk_norm is not None else None, + qk_norm=qk_norm, + dropout=dropout, + bias=attention_bias, + cross_attention_dim=None, + processor=SanaLinearAttnProcessor3_0(), + ) + + # 2. Cross Attention + if cross_attention_dim is not None: + self.norm2 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + self.attn2 = Attention( + query_dim=dim, + qk_norm=qk_norm, + kv_heads=num_cross_attention_heads if qk_norm is not None else None, + cross_attention_dim=cross_attention_dim, + heads=num_cross_attention_heads, + dim_head=cross_attention_head_dim, + dropout=dropout, + bias=True, + out_bias=attention_out_bias, + processor=SanaAttnProcessor2_0(), + ) + + # 3. Feed-forward + self.ff = GLUMBTempConv(dim, dim, mlp_ratio, norm_type=None, residual_connection=False) + + self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + timestep: Optional[torch.LongTensor] = None, + frames: int = None, + height: int = None, + width: int = None, + rotary_emb: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + batch_size = hidden_states.shape[0] + + # 1. Modulation + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1) + ).chunk(6, dim=1) + + # 2. Self Attention + norm_hidden_states = self.norm1(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa + norm_hidden_states = norm_hidden_states.to(hidden_states.dtype) + + attn_output = self.attn1(norm_hidden_states, rotary_emb=rotary_emb) + hidden_states = hidden_states + gate_msa * attn_output + + # 3. Cross Attention + if self.attn2 is not None: + attn_output = self.attn2( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + ) + hidden_states = attn_output + hidden_states + + # 4. Feed-forward + norm_hidden_states = self.norm2(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp + + norm_hidden_states = norm_hidden_states.unflatten(1, (frames, height, width)) + ff_output = self.ff(norm_hidden_states) + ff_output = ff_output.flatten(1, 3) + hidden_states = hidden_states + gate_mlp * ff_output + + return hidden_states + + +class SanaVideoTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, AttentionMixin): + r""" + A 3D Transformer model introduced in [Sana-Video](https://huggingface.co/papers/2509.24695) family of models. + + Args: + in_channels (`int`, defaults to `16`): + The number of channels in the input. + out_channels (`int`, *optional*, defaults to `16`): + The number of channels in the output. + num_attention_heads (`int`, defaults to `20`): + The number of heads to use for multi-head attention. + attention_head_dim (`int`, defaults to `112`): + The number of channels in each head. + num_layers (`int`, defaults to `20`): + The number of layers of Transformer blocks to use. + num_cross_attention_heads (`int`, *optional*, defaults to `20`): + The number of heads to use for cross-attention. + cross_attention_head_dim (`int`, *optional*, defaults to `112`): + The number of channels in each head for cross-attention. + cross_attention_dim (`int`, *optional*, defaults to `2240`): + The number of channels in the cross-attention output. + caption_channels (`int`, defaults to `2304`): + The number of channels in the caption embeddings. + mlp_ratio (`float`, defaults to `2.5`): + The expansion ratio to use in the GLUMBConv layer. + dropout (`float`, defaults to `0.0`): + The dropout probability. + attention_bias (`bool`, defaults to `False`): + Whether to use bias in the attention layer. + sample_size (`int`, defaults to `32`): + The base size of the input latent. + patch_size (`int`, defaults to `1`): + The size of the patches to use in the patch embedding layer. + norm_elementwise_affine (`bool`, defaults to `False`): + Whether to use elementwise affinity in the normalization layer. + norm_eps (`float`, defaults to `1e-6`): + The epsilon value for the normalization layer. + qk_norm (`str`, *optional*, defaults to `None`): + The normalization to use for the query and key. + """ + + _supports_gradient_checkpointing = True + _no_split_modules = ["SanaVideoTransformerBlock", "SanaModulatedNorm"] + _skip_layerwise_casting_patterns = ["patch_embedding", "norm"] + + @register_to_config + def __init__( + self, + in_channels: int = 16, + out_channels: Optional[int] = 16, + num_attention_heads: int = 20, + attention_head_dim: int = 112, + num_layers: int = 20, + num_cross_attention_heads: Optional[int] = 20, + cross_attention_head_dim: Optional[int] = 112, + cross_attention_dim: Optional[int] = 2240, + caption_channels: int = 2304, + mlp_ratio: float = 2.5, + dropout: float = 0.0, + attention_bias: bool = False, + sample_size: int = 30, + patch_size: Tuple[int, int, int] = (1, 2, 2), + norm_elementwise_affine: bool = False, + norm_eps: float = 1e-6, + interpolation_scale: Optional[int] = None, + guidance_embeds: bool = False, + guidance_embeds_scale: float = 0.1, + qk_norm: Optional[str] = "rms_norm_across_heads", + rope_max_seq_len: int = 1024, + ) -> None: + super().__init__() + + out_channels = out_channels or in_channels + inner_dim = num_attention_heads * attention_head_dim + + # 1. Patch & position embedding + self.rope = WanRotaryPosEmbed(attention_head_dim, patch_size, rope_max_seq_len) + self.patch_embedding = nn.Conv3d(in_channels, inner_dim, kernel_size=patch_size, stride=patch_size) + + # 2. Additional condition embeddings + if guidance_embeds: + self.time_embed = SanaCombinedTimestepGuidanceEmbeddings(inner_dim) + else: + self.time_embed = AdaLayerNormSingle(inner_dim) + + self.caption_projection = PixArtAlphaTextProjection(in_features=caption_channels, hidden_size=inner_dim) + self.caption_norm = RMSNorm(inner_dim, eps=1e-5, elementwise_affine=True) + + # 3. Transformer blocks + self.transformer_blocks = nn.ModuleList( + [ + SanaVideoTransformerBlock( + inner_dim, + num_attention_heads, + attention_head_dim, + dropout=dropout, + num_cross_attention_heads=num_cross_attention_heads, + cross_attention_head_dim=cross_attention_head_dim, + cross_attention_dim=cross_attention_dim, + attention_bias=attention_bias, + norm_elementwise_affine=norm_elementwise_affine, + norm_eps=norm_eps, + mlp_ratio=mlp_ratio, + qk_norm=qk_norm, + ) + for _ in range(num_layers) + ] + ) + + # 4. Output blocks + self.scale_shift_table = nn.Parameter(torch.randn(2, inner_dim) / inner_dim**0.5) + self.norm_out = SanaModulatedNorm(inner_dim, elementwise_affine=False, eps=1e-6) + self.proj_out = nn.Linear(inner_dim, math.prod(patch_size) * out_channels) + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.Tensor, + guidance: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + attention_kwargs: Optional[Dict[str, Any]] = None, + controlnet_block_samples: Optional[Tuple[torch.Tensor]] = None, + return_dict: bool = True, + ) -> Union[Tuple[torch.Tensor, ...], Transformer2DModelOutput]: + if attention_kwargs is not None: + attention_kwargs = attention_kwargs.copy() + lora_scale = attention_kwargs.pop("scale", 1.0) + else: + lora_scale = 1.0 + + if USE_PEFT_BACKEND: + # weight the lora layers by setting `lora_scale` for each PEFT layer + scale_lora_layers(self, lora_scale) + else: + if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: + logger.warning( + "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective." + ) + + # ensure attention_mask is a bias, and give it a singleton query_tokens dimension. + # we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward. + # we can tell by counting dims; if ndim == 2: it's a mask rather than a bias. + # expects mask of shape: + # [batch, key_tokens] + # adds singleton query_tokens dimension: + # [batch, 1, key_tokens] + # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes: + # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn) + # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn) + if attention_mask is not None and attention_mask.ndim == 2: + # assume that mask is expressed as: + # (1 = keep, 0 = discard) + # convert mask into a bias that can be added to attention scores: + # (keep = +0, discard = -10000.0) + attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0 + attention_mask = attention_mask.unsqueeze(1) + + # convert encoder_attention_mask to a bias the same way we do for attention_mask + if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2: + encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0 + encoder_attention_mask = encoder_attention_mask.unsqueeze(1) + + # 1. Input + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p_t, p_h, p_w = self.config.patch_size + post_patch_num_frames = num_frames // p_t + post_patch_height = height // p_h + post_patch_width = width // p_w + + rotary_emb = self.rope(hidden_states) + + hidden_states = self.patch_embedding(hidden_states) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + + if guidance is not None: + timestep, embedded_timestep = self.time_embed( + timestep, guidance=guidance, hidden_dtype=hidden_states.dtype + ) + else: + timestep, embedded_timestep = self.time_embed( + timestep, batch_size=batch_size, hidden_dtype=hidden_states.dtype + ) + + encoder_hidden_states = self.caption_projection(encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.shape[-1]) + + encoder_hidden_states = self.caption_norm(encoder_hidden_states) + + # 2. Transformer blocks + if torch.is_grad_enabled() and self.gradient_checkpointing: + for index_block, block in enumerate(self.transformer_blocks): + hidden_states = self._gradient_checkpointing_func( + block, + hidden_states, + attention_mask, + encoder_hidden_states, + encoder_attention_mask, + timestep, + post_patch_num_frames, + post_patch_height, + post_patch_width, + rotary_emb, + ) + if controlnet_block_samples is not None and 0 < index_block <= len(controlnet_block_samples): + hidden_states = hidden_states + controlnet_block_samples[index_block - 1] + + else: + for index_block, block in enumerate(self.transformer_blocks): + hidden_states = block( + hidden_states, + attention_mask, + encoder_hidden_states, + encoder_attention_mask, + timestep, + post_patch_num_frames, + post_patch_height, + post_patch_width, + rotary_emb, + ) + if controlnet_block_samples is not None and 0 < index_block <= len(controlnet_block_samples): + hidden_states = hidden_states + controlnet_block_samples[index_block - 1] + + # 3. Normalization + hidden_states = self.norm_out(hidden_states, embedded_timestep, self.scale_shift_table) + + hidden_states = self.proj_out(hidden_states) + + # 5. Unpatchify + hidden_states = hidden_states.reshape( + batch_size, post_patch_num_frames, post_patch_height, post_patch_width, p_t, p_h, p_w, -1 + ) + hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) + output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + if USE_PEFT_BACKEND: + # remove `lora_scale` from each PEFT layer + unscale_lora_layers(self, lora_scale) + + if not return_dict: + return (output,) + + return Transformer2DModelOutput(sample=output) diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index db357669b6..87d953845e 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -308,6 +308,7 @@ else: "SanaSprintPipeline", "SanaControlNetPipeline", "SanaSprintImg2ImgPipeline", + "SanaVideoPipeline", ] _import_structure["semantic_stable_diffusion"] = ["SemanticStableDiffusionPipeline"] _import_structure["shap_e"] = ["ShapEImg2ImgPipeline", "ShapEPipeline"] @@ -735,7 +736,13 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: QwenImageInpaintPipeline, QwenImagePipeline, ) - from .sana import SanaControlNetPipeline, SanaPipeline, SanaSprintImg2ImgPipeline, SanaSprintPipeline + from .sana import ( + SanaControlNetPipeline, + SanaPipeline, + SanaSprintImg2ImgPipeline, + SanaSprintPipeline, + SanaVideoPipeline, + ) from .semantic_stable_diffusion import SemanticStableDiffusionPipeline from .shap_e import ShapEImg2ImgPipeline, ShapEPipeline from .stable_audio import StableAudioPipeline, StableAudioProjectionModel diff --git a/src/diffusers/pipelines/sana/__init__.py b/src/diffusers/pipelines/sana/__init__.py index 91684f35f1..d5571ab12f 100644 --- a/src/diffusers/pipelines/sana/__init__.py +++ b/src/diffusers/pipelines/sana/__init__.py @@ -26,6 +26,7 @@ else: _import_structure["pipeline_sana_controlnet"] = ["SanaControlNetPipeline"] _import_structure["pipeline_sana_sprint"] = ["SanaSprintPipeline"] _import_structure["pipeline_sana_sprint_img2img"] = ["SanaSprintImg2ImgPipeline"] + _import_structure["pipeline_sana_video"] = ["SanaVideoPipeline"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: try: @@ -39,6 +40,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: from .pipeline_sana_controlnet import SanaControlNetPipeline from .pipeline_sana_sprint import SanaSprintPipeline from .pipeline_sana_sprint_img2img import SanaSprintImg2ImgPipeline + from .pipeline_sana_video import SanaVideoPipeline else: import sys diff --git a/src/diffusers/pipelines/sana/pipeline_output.py b/src/diffusers/pipelines/sana/pipeline_output.py index f8ac129516..8021b77387 100644 --- a/src/diffusers/pipelines/sana/pipeline_output.py +++ b/src/diffusers/pipelines/sana/pipeline_output.py @@ -3,6 +3,7 @@ from typing import List, Union import numpy as np import PIL.Image +import torch from ...utils import BaseOutput @@ -19,3 +20,18 @@ class SanaPipelineOutput(BaseOutput): """ images: Union[List[PIL.Image.Image], np.ndarray] + + +@dataclass +class SanaVideoPipelineOutput(BaseOutput): + r""" + Output class for Sana-Video pipelines. + + Args: + frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]): + List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing + denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape + `(batch_size, num_frames, channels, height, width)`. + """ + + frames: torch.Tensor diff --git a/src/diffusers/pipelines/sana/pipeline_sana.py b/src/diffusers/pipelines/sana/pipeline_sana.py index ac979305ca..2beff802c6 100644 --- a/src/diffusers/pipelines/sana/pipeline_sana.py +++ b/src/diffusers/pipelines/sana/pipeline_sana.py @@ -1,4 +1,4 @@ -# Copyright 2025 PixArt-Sigma Authors and The HuggingFace Team. All rights reserved. +# Copyright 2025 SANA Authors and The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/src/diffusers/pipelines/sana/pipeline_sana_sprint.py b/src/diffusers/pipelines/sana/pipeline_sana_sprint.py index 62b9788292..04f45f817e 100644 --- a/src/diffusers/pipelines/sana/pipeline_sana_sprint.py +++ b/src/diffusers/pipelines/sana/pipeline_sana_sprint.py @@ -1,4 +1,4 @@ -# Copyright 2025 PixArt-Sigma Authors and The HuggingFace Team. All rights reserved. +# Copyright 2025 SANA-Sprint Authors and The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/src/diffusers/pipelines/sana/pipeline_sana_video.py b/src/diffusers/pipelines/sana/pipeline_sana_video.py new file mode 100644 index 0000000000..5ec498faff --- /dev/null +++ b/src/diffusers/pipelines/sana/pipeline_sana_video.py @@ -0,0 +1,1017 @@ +# Copyright 2025 SANA-Video Authors and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import html +import inspect +import re +import urllib.parse as ul +import warnings +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +import torch +from transformers import Gemma2PreTrainedModel, GemmaTokenizer, GemmaTokenizerFast + +from ...callbacks import MultiPipelineCallbacks, PipelineCallback +from ...loaders import SanaLoraLoaderMixin +from ...models import AutoencoderDC, AutoencoderKLWan, SanaVideoTransformer3DModel +from ...schedulers import DPMSolverMultistepScheduler +from ...utils import ( + BACKENDS_MAPPING, + USE_PEFT_BACKEND, + is_bs4_available, + is_ftfy_available, + is_torch_xla_available, + logging, + replace_example_docstring, + scale_lora_layers, + unscale_lora_layers, +) +from ...utils.torch_utils import get_device, is_torch_version, randn_tensor +from ...video_processor import VideoProcessor +from ..pipeline_utils import DiffusionPipeline +from .pipeline_output import SanaVideoPipelineOutput + + +ASPECT_RATIO_480_BIN = { + "0.5": [448.0, 896.0], + "0.57": [480.0, 832.0], + "0.68": [528.0, 768.0], + "0.78": [560.0, 720.0], + "1.0": [624.0, 624.0], + "1.13": [672.0, 592.0], + "1.29": [720.0, 560.0], + "1.46": [768.0, 528.0], + "1.67": [816.0, 496.0], + "1.75": [832.0, 480.0], + "2.0": [896.0, 448.0], +} + + +ASPECT_RATIO_720_BIN = { + "0.5": [672.0, 1344.0], + "0.57": [704.0, 1280.0], + "0.68": [800.0, 1152.0], + "0.78": [832.0, 1088.0], + "1.0": [960.0, 960.0], + "1.13": [1024.0, 896.0], + "1.29": [1088.0, 832.0], + "1.46": [1152.0, 800.0], + "1.67": [1248.0, 736.0], + "1.75": [1280.0, 704.0], + "2.0": [1344.0, 672.0], +} + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +if is_bs4_available(): + from bs4 import BeautifulSoup + +if is_ftfy_available(): + import ftfy + + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import SanaVideoPipeline + >>> from diffusers.utils import export_to_video + + >>> model_id = "Efficient-Large-Model/SANA-Video_2B_480p_diffusers" + >>> pipe = SanaVideoPipeline.from_pretrained(model_id) + >>> pipe.transformer.to(torch.bfloat16) + >>> pipe.text_encoder.to(torch.bfloat16) + >>> pipe.vae.to(torch.float32) + >>> pipe.to("cuda") + >>> model_score = 30 + + >>> prompt = "Evening, backlight, side lighting, soft light, high contrast, mid-shot, centered composition, clean solo shot, warm color. A young Caucasian man stands in a forest, golden light glimmers on his hair as sunlight filters through the leaves. He wears a light shirt, wind gently blowing his hair and collar, light dances across his face with his movements. The background is blurred, with dappled light and soft tree shadows in the distance. The camera focuses on his lifted gaze, clear and emotional." + >>> negative_prompt = "A chaotic sequence with misshapen, deformed limbs in heavy motion blur, sudden disappearance, jump cuts, jerky movements, rapid shot changes, frames out of sync, inconsistent character shapes, temporal artifacts, jitter, and ghosting effects, creating a disorienting visual experience." + >>> motion_prompt = f" motion score: {model_score}." + >>> prompt = prompt + motion_prompt + + >>> output = pipe( + ... prompt=prompt, + ... negative_prompt=negative_prompt, + ... height=480, + ... width=832, + ... frames=81, + ... guidance_scale=6, + ... num_inference_steps=50, + ... generator=torch.Generator(device="cuda").manual_seed(42), + ... ).frames[0] + + >>> export_to_video(output, "sana-video-output.mp4", fps=16) + ``` +""" + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + timesteps: Optional[List[int]] = None, + sigmas: Optional[List[float]] = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`List[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`List[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class SanaVideoPipeline(DiffusionPipeline, SanaLoraLoaderMixin): + r""" + Pipeline for text-to-video generation using [Sana](https://huggingface.co/papers/2509.24695). This model inherits + from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods implemented for all + pipelines (downloading, saving, running on a particular device, etc.). + + Args: + tokenizer ([`GemmaTokenizer`] or [`GemmaTokenizerFast`]): + The tokenizer used to tokenize the prompt. + text_encoder ([`Gemma2PreTrainedModel`]): + Text encoder model to encode the input prompts. + vae ([`AutoencoderKLWan` or `AutoencoderDCAEV`]): + Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations. + transformer ([`SanaVideoTransformer3DModel`]): + Conditional Transformer to denoise the input latents. + scheduler ([`DPMSolverMultistepScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded video latents. + """ + + # fmt: off + bad_punct_regex = re.compile(r"[" + "#®•©™&@·º½¾¿¡§~" + r"\)" + r"\(" + r"\]" + r"\[" + r"\}" + r"\{" + r"\|" + "\\" + r"\/" + r"\*" + r"]{1,}") + # fmt: on + + model_cpu_offload_seq = "text_encoder->transformer->vae" + _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] + + def __init__( + self, + tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast], + text_encoder: Gemma2PreTrainedModel, + vae: Union[AutoencoderDC, AutoencoderKLWan], + transformer: SanaVideoTransformer3DModel, + scheduler: DPMSolverMultistepScheduler, + ): + super().__init__() + + self.register_modules( + tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler + ) + + self.vae_scale_factor_temporal = self.vae.config.scale_factor_temporal if getattr(self, "vae", None) else 4 + self.vae_scale_factor_spatial = self.vae.config.scale_factor_spatial if getattr(self, "vae", None) else 8 + + self.vae_scale_factor = self.vae_scale_factor_spatial + + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) + + def _get_gemma_prompt_embeds( + self, + prompt: Union[str, List[str]], + device: torch.device, + dtype: torch.dtype, + clean_caption: bool = False, + max_sequence_length: int = 300, + complex_human_instruction: Optional[List[str]] = None, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `List[str]`, *optional*): + prompt to be encoded + device: (`torch.device`, *optional*): + torch device to place the resulting embeddings on + clean_caption (`bool`, defaults to `False`): + If `True`, the function will preprocess and clean the provided caption before encoding. + max_sequence_length (`int`, defaults to 300): Maximum sequence length to use for the prompt. + complex_human_instruction (`list[str]`, defaults to `complex_human_instruction`): + If `complex_human_instruction` is not empty, the function will use the complex Human instruction for + the prompt. + """ + prompt = [prompt] if isinstance(prompt, str) else prompt + + if getattr(self, "tokenizer", None) is not None: + self.tokenizer.padding_side = "right" + + prompt = self._text_preprocessing(prompt, clean_caption=clean_caption) + + # prepare complex human instruction + if not complex_human_instruction: + max_length_all = max_sequence_length + else: + chi_prompt = "\n".join(complex_human_instruction) + prompt = [chi_prompt + p for p in prompt] + num_chi_prompt_tokens = len(self.tokenizer.encode(chi_prompt)) + max_length_all = num_chi_prompt_tokens + max_sequence_length - 2 + + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=max_length_all, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + + prompt_attention_mask = text_inputs.attention_mask + prompt_attention_mask = prompt_attention_mask.to(device) + + prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask) + prompt_embeds = prompt_embeds[0].to(dtype=dtype, device=device) + + return prompt_embeds, prompt_attention_mask + + def encode_prompt( + self, + prompt: Union[str, List[str]], + do_classifier_free_guidance: bool = True, + negative_prompt: str = "", + num_videos_per_prompt: int = 1, + device: Optional[torch.device] = None, + prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_attention_mask: Optional[torch.Tensor] = None, + negative_prompt_attention_mask: Optional[torch.Tensor] = None, + clean_caption: bool = False, + max_sequence_length: int = 300, + complex_human_instruction: Optional[List[str]] = None, + lora_scale: Optional[float] = None, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `List[str]`, *optional*): + prompt to be encoded + negative_prompt (`str` or `List[str]`, *optional*): + The prompt not to guide the video 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`). For + PixArt-Alpha, this should be "". + do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): + whether to use classifier free guidance or not + num_videos_per_prompt (`int`, *optional*, defaults to 1): + number of videos that should be generated per prompt + device: (`torch.device`, *optional*): + torch device to place the resulting embeddings on + 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. For Sana, it's should be the embeddings of the "" string. + clean_caption (`bool`, defaults to `False`): + If `True`, the function will preprocess and clean the provided caption before encoding. + max_sequence_length (`int`, defaults to 300): Maximum sequence length to use for the prompt. + complex_human_instruction (`list[str]`, defaults to `complex_human_instruction`): + If `complex_human_instruction` is not empty, the function will use the complex Human instruction for + the prompt. + """ + + if device is None: + device = self._execution_device + + if self.text_encoder is not None: + dtype = self.text_encoder.dtype + else: + dtype = None + + # 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, SanaLoraLoaderMixin): + self._lora_scale = lora_scale + + # dynamically adjust the LoRA scale + if self.text_encoder is not None and USE_PEFT_BACKEND: + scale_lora_layers(self.text_encoder, lora_scale) + + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + if getattr(self, "tokenizer", None) is not None: + self.tokenizer.padding_side = "right" + + # See Section 3.1. of the paper. + max_length = max_sequence_length + select_index = [0] + list(range(-max_length + 1, 0)) + + if prompt_embeds is None: + prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds( + prompt=prompt, + device=device, + dtype=dtype, + clean_caption=clean_caption, + max_sequence_length=max_sequence_length, + complex_human_instruction=complex_human_instruction, + ) + + prompt_embeds = prompt_embeds[:, select_index] + prompt_attention_mask = prompt_attention_mask[:, select_index] + + bs_embed, seq_len, _ = prompt_embeds.shape + # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) + prompt_embeds = prompt_embeds.view(bs_embed * num_videos_per_prompt, seq_len, -1) + prompt_attention_mask = prompt_attention_mask.view(bs_embed, -1) + prompt_attention_mask = prompt_attention_mask.repeat(num_videos_per_prompt, 1) + + # get unconditional embeddings for classifier free guidance + if do_classifier_free_guidance and negative_prompt_embeds is None: + negative_prompt = [negative_prompt] * batch_size if isinstance(negative_prompt, str) else negative_prompt + negative_prompt_embeds, negative_prompt_attention_mask = self._get_gemma_prompt_embeds( + prompt=negative_prompt, + device=device, + dtype=dtype, + clean_caption=clean_caption, + max_sequence_length=max_sequence_length, + complex_human_instruction=False, + ) + + if do_classifier_free_guidance: + # duplicate unconditional embeddings for each generation per prompt, using mps friendly method + seq_len = negative_prompt_embeds.shape[1] + + negative_prompt_embeds = negative_prompt_embeds.to(dtype=dtype, device=device) + + negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_videos_per_prompt, 1) + negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) + + negative_prompt_attention_mask = negative_prompt_attention_mask.view(bs_embed, -1) + negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_videos_per_prompt, 1) + else: + negative_prompt_embeds = None + negative_prompt_attention_mask = None + + if self.text_encoder is not None: + if isinstance(self, SanaLoraLoaderMixin) and USE_PEFT_BACKEND: + # Retrieve the original scale by scaling back the LoRA layers + unscale_lora_layers(self.text_encoder, lora_scale) + + return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask + + # 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://huggingface.co/papers/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 + + def check_inputs( + self, + prompt, + height, + width, + callback_on_step_end_tensor_inputs=None, + negative_prompt=None, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_attention_mask=None, + negative_prompt_attention_mask=None, + ): + if height % 32 != 0 or width % 32 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 32 but are {height} and {width}.") + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + if prompt_embeds is not None and prompt_attention_mask is None: + raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.") + + if negative_prompt_embeds is not None and negative_prompt_attention_mask is None: + raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.") + + if prompt_embeds is not None and negative_prompt_embeds is not None: + if prompt_embeds.shape != negative_prompt_embeds.shape: + raise ValueError( + "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" + f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" + f" {negative_prompt_embeds.shape}." + ) + if prompt_attention_mask.shape != negative_prompt_attention_mask.shape: + raise ValueError( + "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but" + f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`" + f" {negative_prompt_attention_mask.shape}." + ) + + # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing + def _text_preprocessing(self, text, clean_caption=False): + if clean_caption and not is_bs4_available(): + logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`")) + logger.warning("Setting `clean_caption` to False...") + clean_caption = False + + if clean_caption and not is_ftfy_available(): + logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`")) + logger.warning("Setting `clean_caption` to False...") + clean_caption = False + + if not isinstance(text, (tuple, list)): + text = [text] + + def process(text: str): + if clean_caption: + text = self._clean_caption(text) + text = self._clean_caption(text) + else: + text = text.lower().strip() + return text + + return [process(t) for t in text] + + # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption + def _clean_caption(self, caption): + caption = str(caption) + caption = ul.unquote_plus(caption) + caption = caption.strip().lower() + caption = re.sub("", "person", caption) + # urls: + caption = re.sub( + r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa + "", + caption, + ) # regex for urls + caption = re.sub( + r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa + "", + caption, + ) # regex for urls + # html: + caption = BeautifulSoup(caption, features="html.parser").text + + # @ + caption = re.sub(r"@[\w\d]+\b", "", caption) + + # 31C0—31EF CJK Strokes + # 31F0—31FF Katakana Phonetic Extensions + # 3200—32FF Enclosed CJK Letters and Months + # 3300—33FF CJK Compatibility + # 3400—4DBF CJK Unified Ideographs Extension A + # 4DC0—4DFF Yijing Hexagram Symbols + # 4E00—9FFF CJK Unified Ideographs + caption = re.sub(r"[\u31c0-\u31ef]+", "", caption) + caption = re.sub(r"[\u31f0-\u31ff]+", "", caption) + caption = re.sub(r"[\u3200-\u32ff]+", "", caption) + caption = re.sub(r"[\u3300-\u33ff]+", "", caption) + caption = re.sub(r"[\u3400-\u4dbf]+", "", caption) + caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption) + caption = re.sub(r"[\u4e00-\u9fff]+", "", caption) + ####################################################### + + # все виды тире / all types of dash --> "-" + caption = re.sub( + r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+", # noqa + "-", + caption, + ) + + # кавычки к одному стандарту + caption = re.sub(r"[`´«»“”¨]", '"', caption) + caption = re.sub(r"[‘’]", "'", caption) + + # " + caption = re.sub(r""?", "", caption) + # & + caption = re.sub(r"&", "", caption) + + # ip addresses: + caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption) + + # article ids: + caption = re.sub(r"\d:\d\d\s+$", "", caption) + + # \n + caption = re.sub(r"\\n", " ", caption) + + # "#123" + caption = re.sub(r"#\d{1,3}\b", "", caption) + # "#12345.." + caption = re.sub(r"#\d{5,}\b", "", caption) + # "123456.." + caption = re.sub(r"\b\d{6,}\b", "", caption) + # filenames: + caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption) + + # + caption = re.sub(r"[\"\']{2,}", r'"', caption) # """AUSVERKAUFT""" + caption = re.sub(r"[\.]{2,}", r" ", caption) # """AUSVERKAUFT""" + + caption = re.sub(self.bad_punct_regex, r" ", caption) # ***AUSVERKAUFT***, #AUSVERKAUFT + caption = re.sub(r"\s+\.\s+", r" ", caption) # " . " + + # this-is-my-cute-cat / this_is_my_cute_cat + regex2 = re.compile(r"(?:\-|\_)") + if len(re.findall(regex2, caption)) > 3: + caption = re.sub(regex2, " ", caption) + + caption = ftfy.fix_text(caption) + caption = html.unescape(html.unescape(caption)) + + caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption) # jc6640 + caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption) # jc6640vc + caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption) # 6640vc231 + + caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption) + caption = re.sub(r"(free\s)?download(\sfree)?", "", caption) + caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption) + caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption) + caption = re.sub(r"\bpage\s+\d+\b", "", caption) + + caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption) # j2d1a2a... + + caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption) + + caption = re.sub(r"\b\s+\:\s+", r": ", caption) + caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption) + caption = re.sub(r"\s+", " ", caption) + + caption.strip() + + caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption) + caption = re.sub(r"^[\'\_,\-\:;]", r"", caption) + caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption) + caption = re.sub(r"^\.\S+$", "", caption) + + return caption.strip() + + def prepare_latents( + self, + batch_size: int, + num_channels_latents: int = 16, + height: int = 480, + width: int = 832, + num_frames: int = 81, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=device, dtype=dtype) + + num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + shape = ( + batch_size, + num_channels_latents, + num_latent_frames, + int(height) // self.vae_scale_factor_spatial, + int(width) // self.vae_scale_factor_spatial, + ) + 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=device, dtype=dtype) + return latents + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def attention_kwargs(self): + return self._attention_kwargs + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1.0 + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: Union[str, List[str]] = None, + negative_prompt: str = "", + num_inference_steps: int = 50, + timesteps: List[int] = None, + sigmas: List[float] = None, + guidance_scale: float = 6.0, + num_videos_per_prompt: Optional[int] = 1, + height: int = 480, + width: int = 832, + frames: int = 81, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + prompt_embeds: Optional[torch.Tensor] = None, + prompt_attention_mask: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt_attention_mask: Optional[torch.Tensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + clean_caption: bool = False, + use_resolution_binning: bool = True, + attention_kwargs: Optional[Dict[str, Any]] = None, + callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + max_sequence_length: int = 300, + complex_human_instruction: List[str] = [ + "Given a user prompt, generate an 'Enhanced prompt' that provides detailed visual descriptions suitable for video generation. Evaluate the level of detail in the user prompt:", + "- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, motion, and temporal relationships to create vivid and dynamic scenes.", + "- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.", + "Here are examples of how to transform or refine prompts:", + "- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat slowly settling into a curled position, peacefully falling asleep on a warm sunny windowsill, with gentle sunlight filtering through surrounding pots of blooming red flowers.", + "- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps gradually lighting up, a diverse crowd of people in colorful clothing walking past, and a double-decker bus smoothly passing by towering glass skyscrapers.", + "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:", + "User Prompt: ", + ], + ) -> Union[SanaVideoPipelineOutput, Tuple]: + """ + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the video generation. If not defined, one has to pass `prompt_embeds`. + instead. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the video generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality video at the + expense of slower inference. + timesteps (`List[int]`, *optional*): + 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. + sigmas (`List[float]`, *optional*): + Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in + their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed + will be used. + guidance_scale (`float`, *optional*, defaults to 4.5): + Guidance scale as defined in [Classifier-Free Diffusion + Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. + of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting + `guidance_scale > 1`. Higher guidance scale encourages to generate videos that are closely linked to + the text `prompt`, usually at the expense of lower video quality. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of videos to generate per prompt. + height (`int`, *optional*, defaults to 480): + The height in pixels of the generated video. + width (`int`, *optional*, defaults to 832): + The width in pixels of the generated video. + frames (`int`, *optional*, defaults to 81): + The number of frames in the generated video. + eta (`float`, *optional*, defaults to 0.0): + Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only + applies to [`schedulers.DDIMScheduler`], will be ignored for others. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for video + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will be generated by sampling using the supplied random `generator`. + prompt_embeds (`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. + prompt_attention_mask (`torch.Tensor`, *optional*): Pre-generated attention mask for text embeddings. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. For PixArt-Sigma this negative prompt should be "". If not + provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. + negative_prompt_attention_mask (`torch.Tensor`, *optional*): + Pre-generated attention mask for negative text embeddings. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generated video. Choose between mp4 or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`SanaVideoPipelineOutput`] instead of a plain tuple. + attention_kwargs: + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + clean_caption (`bool`, *optional*, defaults to `True`): + Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to + be installed. If the dependencies are not installed, the embeddings will be created from the raw + prompt. + use_resolution_binning (`bool` defaults to `True`): + If set to `True`, the requested height and width are first mapped to the closest resolutions using + `ASPECT_RATIO_480_BIN` or `ASPECT_RATIO_720_BIN`. After the produced latents are decoded into videos, + they are resized back to the requested resolution. Useful for generating non-square videos. + callback_on_step_end (`Callable`, *optional*): + A function that calls at the end of each denoising steps during the inference. The function is called + with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, + callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by + `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + max_sequence_length (`int` defaults to `300`): + Maximum sequence length to use with the `prompt`. + complex_human_instruction (`List[str]`, *optional*): + Instructions for complex human attention: + https://github.com/NVlabs/Sana/blob/main/configs/sana_app_config/Sana_1600M_app.yaml#L55. + + Examples: + + Returns: + [`~pipelines.sana.pipeline_output.SanaVideoPipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`~pipelines.sana.pipeline_output.SanaVideoPipelineOutput`] is returned, + otherwise a `tuple` is returned where the first element is a list with the generated videos + """ + + if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): + callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs + + # 1. Check inputs. Raise error if not correct + if use_resolution_binning: + if self.transformer.config.sample_size == 30: + aspect_ratio_bin = ASPECT_RATIO_480_BIN + elif self.transformer.config.sample_size == 22: + aspect_ratio_bin = ASPECT_RATIO_720_BIN + else: + raise ValueError("Invalid sample size") + orig_height, orig_width = height, width + height, width = self.video_processor.classify_height_width_bin(height, width, ratios=aspect_ratio_bin) + + self.check_inputs( + prompt, + height, + width, + callback_on_step_end_tensor_inputs, + negative_prompt, + prompt_embeds, + negative_prompt_embeds, + prompt_attention_mask, + negative_prompt_attention_mask, + ) + + self._guidance_scale = guidance_scale + self._attention_kwargs = attention_kwargs + self._interrupt = False + + # 2. Default height and width to transformer + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self._execution_device + lora_scale = self.attention_kwargs.get("scale", None) if self.attention_kwargs is not None else None + + # 3. Encode input prompt + ( + prompt_embeds, + prompt_attention_mask, + negative_prompt_embeds, + negative_prompt_attention_mask, + ) = self.encode_prompt( + prompt, + self.do_classifier_free_guidance, + negative_prompt=negative_prompt, + num_videos_per_prompt=num_videos_per_prompt, + device=device, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + clean_caption=clean_caption, + max_sequence_length=max_sequence_length, + complex_human_instruction=complex_human_instruction, + lora_scale=lora_scale, + ) + if self.do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) + + # 4. Prepare timesteps + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, device, timesteps, sigmas + ) + + # 5. Prepare latents. + latent_channels = self.transformer.config.in_channels + latents = self.prepare_latents( + batch_size * num_videos_per_prompt, + latent_channels, + height, + width, + frames, + torch.float32, + device, + generator, + latents, + ) + + # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + + # 7. Denoising loop + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + + transformer_dtype = self.transformer.dtype + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latent_model_input.shape[0]) + + # predict noise model_output + noise_pred = self.transformer( + latent_model_input.to(dtype=transformer_dtype), + encoder_hidden_states=prompt_embeds.to(dtype=transformer_dtype), + encoder_attention_mask=prompt_attention_mask, + timestep=timestep, + return_dict=False, + attention_kwargs=self.attention_kwargs, + )[0] + noise_pred = noise_pred.float() + + # perform guidance + if self.do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + # learned sigma + if self.transformer.config.out_channels // 2 == latent_channels: + noise_pred = noise_pred.chunk(2, dim=1)[0] + + # compute previous image: x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + if output_type == "latent": + video = latents + else: + latents = latents.to(self.vae.dtype) + torch_accelerator_module = getattr(torch, get_device(), torch.cuda) + oom_error = ( + torch.OutOfMemoryError + if is_torch_version(">=", "2.5.0") + else torch_accelerator_module.OutOfMemoryError + ) + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( + latents.device, latents.dtype + ) + latents = latents / latents_std + latents_mean + try: + video = self.vae.decode(latents, return_dict=False)[0] + except oom_error as e: + warnings.warn( + f"{e}. \n" + f"Try to use VAE tiling for large images. For example: \n" + f"pipe.vae.enable_tiling(tile_sample_min_width=512, tile_sample_min_height=512)" + ) + + if use_resolution_binning: + video = self.video_processor.resize_and_crop_tensor(video, orig_width, orig_height) + + video = self.video_processor.postprocess_video(video, output_type=output_type) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (video,) + + return SanaVideoPipelineOutput(frames=video) diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 3c426d5039..22d2d8c0a5 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -1308,6 +1308,21 @@ class SanaTransformer2DModel(metaclass=DummyObject): requires_backends(cls, ["torch"]) +class SanaVideoTransformer3DModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class SD3ControlNetModel(metaclass=DummyObject): _backends = ["torch"] diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 20575ff229..e8209403de 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -2177,6 +2177,21 @@ class SanaSprintPipeline(metaclass=DummyObject): requires_backends(cls, ["torch", "transformers"]) +class SanaVideoPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class SemanticStableDiffusionPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/src/diffusers/video_processor.py b/src/diffusers/video_processor.py index 59b59b47d2..abeb30bca1 100644 --- a/src/diffusers/video_processor.py +++ b/src/diffusers/video_processor.py @@ -13,11 +13,12 @@ # limitations under the License. import warnings -from typing import List, Optional, Union +from typing import List, Optional, Tuple, Union import numpy as np import PIL import torch +import torch.nn.functional as F from .image_processor import VaeImageProcessor, is_valid_image, is_valid_image_imagelist @@ -111,3 +112,65 @@ class VideoProcessor(VaeImageProcessor): raise ValueError(f"{output_type} does not exist. Please choose one of ['np', 'pt', 'pil']") return outputs + + @staticmethod + def classify_height_width_bin(height: int, width: int, ratios: dict) -> Tuple[int, int]: + r""" + Returns the binned height and width based on the aspect ratio. + + Args: + height (`int`): The height of the image. + width (`int`): The width of the image. + ratios (`dict`): A dictionary where keys are aspect ratios and values are tuples of (height, width). + + Returns: + `Tuple[int, int]`: The closest binned height and width. + """ + ar = float(height / width) + closest_ratio = min(ratios.keys(), key=lambda ratio: abs(float(ratio) - ar)) + default_hw = ratios[closest_ratio] + return int(default_hw[0]), int(default_hw[1]) + + @staticmethod + def resize_and_crop_tensor(samples: torch.Tensor, new_width: int, new_height: int) -> torch.Tensor: + r""" + Resizes and crops a tensor of videos to the specified dimensions. + + Args: + samples (`torch.Tensor`): + A tensor of shape (N, C, T, H, W) where N is the batch size, C is the number of channels, T is the + number of frames, H is the height, and W is the width. + new_width (`int`): The desired width of the output videos. + new_height (`int`): The desired height of the output videos. + + Returns: + `torch.Tensor`: A tensor containing the resized and cropped videos. + """ + orig_height, orig_width = samples.shape[3], samples.shape[4] + + # Check if resizing is needed + if orig_height != new_height or orig_width != new_width: + ratio = max(new_height / orig_height, new_width / orig_width) + resized_width = int(orig_width * ratio) + resized_height = int(orig_height * ratio) + + # Reshape to (N*T, C, H, W) for interpolation + n, c, t, h, w = samples.shape + samples = samples.permute(0, 2, 1, 3, 4).reshape(n * t, c, h, w) + + # Resize + samples = F.interpolate( + samples, size=(resized_height, resized_width), mode="bilinear", align_corners=False + ) + + # Center Crop + start_x = (resized_width - new_width) // 2 + end_x = start_x + new_width + start_y = (resized_height - new_height) // 2 + end_y = start_y + new_height + samples = samples[:, :, start_y:end_y, start_x:end_x] + + # Reshape back to (N, C, T, H, W) + samples = samples.reshape(n, t, c, new_height, new_width).permute(0, 2, 1, 3, 4) + + return samples diff --git a/tests/models/transformers/test_models_transformer_sana_video.py b/tests/models/transformers/test_models_transformer_sana_video.py new file mode 100644 index 0000000000..ff564ed891 --- /dev/null +++ b/tests/models/transformers/test_models_transformer_sana_video.py @@ -0,0 +1,97 @@ +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import torch + +from diffusers import SanaVideoTransformer3DModel + +from ...testing_utils import ( + enable_full_determinism, + torch_device, +) +from ..test_modeling_common import ModelTesterMixin, TorchCompileTesterMixin + + +enable_full_determinism() + + +class SanaVideoTransformer3DTests(ModelTesterMixin, unittest.TestCase): + model_class = SanaVideoTransformer3DModel + main_input_name = "hidden_states" + uses_custom_attn_processor = True + + @property + def dummy_input(self): + batch_size = 1 + num_channels = 16 + num_frames = 2 + height = 16 + width = 16 + text_encoder_embedding_dim = 16 + sequence_length = 12 + + hidden_states = torch.randn((batch_size, num_channels, num_frames, height, width)).to(torch_device) + timestep = torch.randint(0, 1000, size=(batch_size,)).to(torch_device) + encoder_hidden_states = torch.randn((batch_size, sequence_length, text_encoder_embedding_dim)).to(torch_device) + + return { + "hidden_states": hidden_states, + "encoder_hidden_states": encoder_hidden_states, + "timestep": timestep, + } + + @property + def input_shape(self): + return (16, 2, 16, 16) + + @property + def output_shape(self): + return (16, 2, 16, 16) + + def prepare_init_args_and_inputs_for_common(self): + init_dict = { + "in_channels": 16, + "out_channels": 16, + "num_attention_heads": 2, + "attention_head_dim": 12, + "num_layers": 2, + "num_cross_attention_heads": 2, + "cross_attention_head_dim": 12, + "cross_attention_dim": 24, + "caption_channels": 16, + "mlp_ratio": 2.5, + "dropout": 0.0, + "attention_bias": False, + "sample_size": 8, + "patch_size": (1, 2, 2), + "norm_elementwise_affine": False, + "norm_eps": 1e-6, + "qk_norm": "rms_norm_across_heads", + "rope_max_seq_len": 32, + } + inputs_dict = self.dummy_input + return init_dict, inputs_dict + + def test_gradient_checkpointing_is_applied(self): + expected_set = {"SanaVideoTransformer3DModel"} + super().test_gradient_checkpointing_is_applied(expected_set=expected_set) + + +class SanaVideoTransformerCompileTests(TorchCompileTesterMixin, unittest.TestCase): + model_class = SanaVideoTransformer3DModel + + def prepare_init_args_and_inputs_for_common(self): + return SanaVideoTransformer3DTests().prepare_init_args_and_inputs_for_common() diff --git a/tests/pipelines/sana/test_sana_video.py b/tests/pipelines/sana/test_sana_video.py new file mode 100644 index 0000000000..9f360a942a --- /dev/null +++ b/tests/pipelines/sana/test_sana_video.py @@ -0,0 +1,225 @@ +# Copyright 2025 The HuggingFace Team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import gc +import tempfile +import unittest + +import numpy as np +import torch +from transformers import Gemma2Config, Gemma2Model, GemmaTokenizer + +from diffusers import AutoencoderKLWan, DPMSolverMultistepScheduler, SanaVideoPipeline, SanaVideoTransformer3DModel + +from ...testing_utils import ( + backend_empty_cache, + enable_full_determinism, + require_torch_accelerator, + slow, + torch_device, +) +from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS +from ..test_pipelines_common import PipelineTesterMixin + + +enable_full_determinism() + + +class SanaVideoPipelineFastTests(PipelineTesterMixin, unittest.TestCase): + pipeline_class = SanaVideoPipeline + params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} + batch_params = TEXT_TO_IMAGE_BATCH_PARAMS + image_params = TEXT_TO_IMAGE_IMAGE_PARAMS + image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS + required_optional_params = frozenset( + [ + "num_inference_steps", + "generator", + "latents", + "return_dict", + "callback_on_step_end", + "callback_on_step_end_tensor_inputs", + ] + ) + test_xformers_attention = False + supports_dduf = False + + def get_dummy_components(self): + torch.manual_seed(0) + vae = AutoencoderKLWan( + base_dim=3, + z_dim=16, + dim_mult=[1, 1, 1, 1], + num_res_blocks=1, + temperal_downsample=[False, True, True], + ) + + torch.manual_seed(0) + scheduler = DPMSolverMultistepScheduler() + + torch.manual_seed(0) + text_encoder_config = Gemma2Config( + head_dim=16, + hidden_size=8, + initializer_range=0.02, + intermediate_size=64, + max_position_embeddings=8192, + model_type="gemma2", + num_attention_heads=2, + num_hidden_layers=1, + num_key_value_heads=2, + vocab_size=8, + attn_implementation="eager", + ) + text_encoder = Gemma2Model(text_encoder_config) + tokenizer = GemmaTokenizer.from_pretrained("hf-internal-testing/dummy-gemma") + + torch.manual_seed(0) + transformer = SanaVideoTransformer3DModel( + in_channels=16, + out_channels=16, + num_attention_heads=2, + attention_head_dim=12, + num_layers=2, + num_cross_attention_heads=2, + cross_attention_head_dim=12, + cross_attention_dim=24, + caption_channels=8, + mlp_ratio=2.5, + dropout=0.0, + attention_bias=False, + sample_size=8, + patch_size=(1, 2, 2), + norm_elementwise_affine=False, + norm_eps=1e-6, + qk_norm="rms_norm_across_heads", + rope_max_seq_len=32, + ) + + components = { + "transformer": transformer, + "vae": vae, + "scheduler": scheduler, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + } + return components + + def get_dummy_inputs(self, device, seed=0): + if str(device).startswith("mps"): + generator = torch.manual_seed(seed) + else: + generator = torch.Generator(device=device).manual_seed(seed) + inputs = { + "prompt": "", + "negative_prompt": "", + "generator": generator, + "num_inference_steps": 2, + "guidance_scale": 6.0, + "height": 32, + "width": 32, + "frames": 9, + "max_sequence_length": 16, + "output_type": "pt", + "complex_human_instruction": [], + "use_resolution_binning": False, + } + return inputs + + def test_inference(self): + device = "cpu" + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + video = pipe(**inputs).frames + generated_video = video[0] + self.assertEqual(generated_video.shape, (9, 3, 32, 32)) + + @unittest.skip("Test not supported") + def test_attention_slicing_forward_pass(self): + pass + + def test_save_load_local(self, expected_max_difference=5e-4): + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + for component in pipe.components.values(): + if hasattr(component, "set_default_attn_processor"): + component.set_default_attn_processor() + pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(torch_device) + torch.manual_seed(0) + output = pipe(**inputs)[0] + + with tempfile.TemporaryDirectory() as tmpdir: + pipe.save_pretrained(tmpdir, safe_serialization=False) + pipe_loaded = self.pipeline_class.from_pretrained(tmpdir) + for component in pipe_loaded.components.values(): + if hasattr(component, "set_default_attn_processor"): + component.set_default_attn_processor() + pipe_loaded.to(torch_device) + pipe_loaded.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(torch_device) + torch.manual_seed(0) + output_loaded = pipe_loaded(**inputs)[0] + + max_diff = np.abs(output.detach().cpu().numpy() - output_loaded.detach().cpu().numpy()).max() + self.assertLess(max_diff, expected_max_difference) + + # TODO(aryan): Create a dummy gemma model with smol vocab size + @unittest.skip( + "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." + ) + def test_inference_batch_consistent(self): + pass + + @unittest.skip( + "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." + ) + def test_inference_batch_single_identical(self): + pass + + def test_float16_inference(self): + # Requires higher tolerance as model seems very sensitive to dtype + super().test_float16_inference(expected_max_diff=0.08) + + def test_save_load_float16(self): + # Requires higher tolerance as model seems very sensitive to dtype + super().test_save_load_float16(expected_max_diff=0.2) + + +@slow +@require_torch_accelerator +class SanaVideoPipelineIntegrationTests(unittest.TestCase): + prompt = "Evening, backlight, side lighting, soft light, high contrast, mid-shot, centered composition, clean solo shot, warm color. A young Caucasian man stands in a forest." + + def setUp(self): + super().setUp() + gc.collect() + backend_empty_cache(torch_device) + + def tearDown(self): + super().tearDown() + gc.collect() + backend_empty_cache(torch_device) + + @unittest.skip("TODO: test needs to be implemented") + def test_sana_video_480p(self): + pass From e4393fa613e0ffbafe3b07a70b05a042f3a26537 Mon Sep 17 00:00:00 2001 From: Mohammad Sadegh Salehi <34940948+MohammadSadeghSalehi@users.noreply.github.com> Date: Thu, 6 Nov 2025 18:18:21 +0000 Subject: [PATCH 122/123] Fix overflow and dtype handling in rgblike_to_depthmap (NumPy + PyTorch) (#12546) * Fix overflow in rgblike_to_depthmap by safe dtype casting (torch & NumPy) * Fix: store original dtype and cast back after safe computation * Apply style fixes --------- Co-authored-by: github-actions[bot] --- src/diffusers/image_processor.py | 41 +++++++++++++++++++++++++------- 1 file changed, 32 insertions(+), 9 deletions(-) diff --git a/src/diffusers/image_processor.py b/src/diffusers/image_processor.py index 0e3082eada..067d876ffc 100644 --- a/src/diffusers/image_processor.py +++ b/src/diffusers/image_processor.py @@ -1045,16 +1045,39 @@ class VaeImageProcessorLDM3D(VaeImageProcessor): def rgblike_to_depthmap(image: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: r""" Convert an RGB-like depth image to a depth map. - - Args: - image (`Union[np.ndarray, torch.Tensor]`): - The RGB-like depth image to convert. - - Returns: - `Union[np.ndarray, torch.Tensor]`: - The corresponding depth map. """ - return image[:, :, 1] * 2**8 + image[:, :, 2] + # 1. Cast the tensor to a larger integer type (e.g., int32) + # to safely perform the multiplication by 256. + # 2. Perform the 16-bit combination: High-byte * 256 + Low-byte. + # 3. Cast the final result to the desired depth map type (uint16) if needed + # before returning, though leaving it as int32/int64 is often safer + # for return value from a library function. + + if isinstance(image, torch.Tensor): + # Cast to a safe dtype (e.g., int32 or int64) for the calculation + original_dtype = image.dtype + image_safe = image.to(torch.int32) + + # Calculate the depth map + depth_map = image_safe[:, :, 1] * 256 + image_safe[:, :, 2] + + # You may want to cast the final result to uint16, but casting to a + # larger int type (like int32) is sufficient to fix the overflow. + # depth_map = depth_map.to(torch.uint16) # Uncomment if uint16 is strictly required + return depth_map.to(original_dtype) + + elif isinstance(image, np.ndarray): + # NumPy equivalent: Cast to a safe dtype (e.g., np.int32) + original_dtype = image.dtype + image_safe = image.astype(np.int32) + + # Calculate the depth map + depth_map = image_safe[:, :, 1] * 256 + image_safe[:, :, 2] + + # depth_map = depth_map.astype(np.uint16) # Uncomment if uint16 is strictly required + return depth_map.astype(original_dtype) + else: + raise TypeError("Input image must be a torch.Tensor or np.ndarray") def numpy_to_depth(self, images: np.ndarray) -> List[PIL.Image.Image]: r""" From 8ac17cd2cbac738af5e0f6c2b3d971dfac552758 Mon Sep 17 00:00:00 2001 From: Dhruv Nair Date: Fri, 7 Nov 2025 08:19:15 +0530 Subject: [PATCH 123/123] [Modular] Some clean up for Modular tests (#12579) * update * update --------- Co-authored-by: Sayak Paul --- .../flux/test_modular_pipeline_flux.py | 110 +++++-- ...st_modular_pipeline_stable_diffusion_xl.py | 295 ++++++++++-------- .../test_modular_pipelines_common.py | 63 ++-- 3 files changed, 269 insertions(+), 199 deletions(-) diff --git a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py index 9d70c21aa8..a29fd43614 100644 --- a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py +++ b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py @@ -15,7 +15,6 @@ import random import tempfile -import unittest import numpy as np import PIL @@ -34,21 +33,16 @@ from ...testing_utils import floats_tensor, torch_device from ..test_modular_pipelines_common import ModularPipelineTesterMixin -class FluxModularTests: +class TestFluxModularPipelineFast(ModularPipelineTesterMixin): pipeline_class = FluxModularPipeline pipeline_blocks_class = FluxAutoBlocks repo = "hf-internal-testing/tiny-flux-modular" - def get_pipeline(self, components_manager=None, torch_dtype=torch.float32): - pipeline = self.pipeline_blocks_class().init_pipeline(self.repo, components_manager=components_manager) - pipeline.load_components(torch_dtype=torch_dtype) - return pipeline + params = frozenset(["prompt", "height", "width", "guidance_scale"]) + batch_params = frozenset(["prompt"]) - def get_dummy_inputs(self, device, seed=0): - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device=device).manual_seed(seed) + def get_dummy_inputs(self, seed=0): + generator = self.get_generator(seed) inputs = { "prompt": "A painting of a squirrel eating a burger", "generator": generator, @@ -57,36 +51,47 @@ class FluxModularTests: "height": 8, "width": 8, "max_sequence_length": 48, - "output_type": "np", + "output_type": "pt", } return inputs -class FluxModularPipelineFastTests(FluxModularTests, ModularPipelineTesterMixin, unittest.TestCase): - params = frozenset(["prompt", "height", "width", "guidance_scale"]) - batch_params = frozenset(["prompt"]) +class TestFluxImg2ImgModularPipelineFast(ModularPipelineTesterMixin): + pipeline_class = FluxModularPipeline + pipeline_blocks_class = FluxAutoBlocks + repo = "hf-internal-testing/tiny-flux-modular" - -class FluxImg2ImgModularPipelineFastTests(FluxModularTests, ModularPipelineTesterMixin, unittest.TestCase): params = frozenset(["prompt", "height", "width", "guidance_scale", "image"]) batch_params = frozenset(["prompt", "image"]) def get_pipeline(self, components_manager=None, torch_dtype=torch.float32): pipeline = super().get_pipeline(components_manager, torch_dtype) + # Override `vae_scale_factor` here as currently, `image_processor` is initialized with # fixed constants instead of # https://github.com/huggingface/diffusers/blob/d54622c2679d700b425ad61abce9b80fc36212c0/src/diffusers/pipelines/flux/pipeline_flux_img2img.py#L230C9-L232C10 pipeline.image_processor = VaeImageProcessor(vae_scale_factor=2) return pipeline - def get_dummy_inputs(self, device, seed=0): - inputs = super().get_dummy_inputs(device, seed) - image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device) - image = image / 2 + 0.5 - inputs["image"] = image - inputs["strength"] = 0.8 - inputs["height"] = 8 - inputs["width"] = 8 + def get_dummy_inputs(self, seed=0): + generator = self.get_generator(seed) + inputs = { + "prompt": "A painting of a squirrel eating a burger", + "generator": generator, + "num_inference_steps": 4, + "guidance_scale": 5.0, + "height": 8, + "width": 8, + "max_sequence_length": 48, + "output_type": "pt", + } + image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(torch_device) + image = image.cpu().permute(0, 2, 3, 1)[0] + init_image = PIL.Image.fromarray(np.uint8(image)).convert("RGB") + + inputs["image"] = init_image + inputs["strength"] = 0.5 + return inputs def test_save_from_pretrained(self): @@ -96,6 +101,7 @@ class FluxImg2ImgModularPipelineFastTests(FluxModularTests, ModularPipelineTeste with tempfile.TemporaryDirectory() as tmpdirname: base_pipe.save_pretrained(tmpdirname) + pipe = ModularPipeline.from_pretrained(tmpdirname).to(torch_device) pipe.load_components(torch_dtype=torch.float32) pipe.to(torch_device) @@ -105,26 +111,62 @@ class FluxImg2ImgModularPipelineFastTests(FluxModularTests, ModularPipelineTeste image_slices = [] for pipe in pipes: - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() image = pipe(**inputs, output="images") image_slices.append(image[0, -3:, -3:, -1].flatten()) - assert np.abs(image_slices[0] - image_slices[1]).max() < 1e-3 + assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3 -class FluxKontextModularPipelineFastTests(FluxImg2ImgModularPipelineFastTests): +class TestFluxKontextModularPipelineFast(ModularPipelineTesterMixin): pipeline_class = FluxKontextModularPipeline pipeline_blocks_class = FluxKontextAutoBlocks repo = "hf-internal-testing/tiny-flux-kontext-pipe" - def get_dummy_inputs(self, device, seed=0): - inputs = super().get_dummy_inputs(device, seed) + params = frozenset(["prompt", "height", "width", "guidance_scale", "image"]) + batch_params = frozenset(["prompt", "image"]) + + def get_dummy_inputs(self, seed=0): + generator = self.get_generator(seed) + inputs = { + "prompt": "A painting of a squirrel eating a burger", + "generator": generator, + "num_inference_steps": 2, + "guidance_scale": 5.0, + "height": 8, + "width": 8, + "max_sequence_length": 48, + "output_type": "pt", + } image = PIL.Image.new("RGB", (32, 32), 0) - _ = inputs.pop("strength") + inputs["image"] = image - inputs["height"] = 8 - inputs["width"] = 8 - inputs["max_area"] = 8 * 8 + inputs["max_area"] = inputs["height"] * inputs["width"] inputs["_auto_resize"] = False + return inputs + + def test_save_from_pretrained(self): + pipes = [] + base_pipe = self.get_pipeline().to(torch_device) + pipes.append(base_pipe) + + with tempfile.TemporaryDirectory() as tmpdirname: + base_pipe.save_pretrained(tmpdirname) + + pipe = ModularPipeline.from_pretrained(tmpdirname).to(torch_device) + pipe.load_components(torch_dtype=torch.float32) + pipe.to(torch_device) + pipe.image_processor = VaeImageProcessor(vae_scale_factor=2) + + pipes.append(pipe) + + image_slices = [] + for pipe in pipes: + inputs = self.get_dummy_inputs() + image = pipe(**inputs, output="images") + + image_slices.append(image[0, -3:, -3:, -1].flatten()) + + assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3 diff --git a/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py b/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py index 22347aa558..ea54b2bdff 100644 --- a/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py +++ b/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py @@ -14,7 +14,6 @@ # limitations under the License. import random -import unittest from typing import Any, Dict import numpy as np @@ -32,63 +31,26 @@ from ..test_modular_pipelines_common import ModularPipelineTesterMixin enable_full_determinism() -class SDXLModularTests: +class SDXLModularTesterMixin: """ This mixin defines method to create pipeline, base input and base test across all SDXL modular tests. """ - pipeline_class = StableDiffusionXLModularPipeline - pipeline_blocks_class = StableDiffusionXLAutoBlocks - repo = "hf-internal-testing/tiny-sdxl-modular" - params = frozenset( - [ - "prompt", - "height", - "width", - "negative_prompt", - "cross_attention_kwargs", - "image", - "mask_image", - ] - ) - batch_params = frozenset(["prompt", "negative_prompt", "image", "mask_image"]) - - def get_pipeline(self, components_manager=None, torch_dtype=torch.float32): - pipeline = self.pipeline_blocks_class().init_pipeline(self.repo, components_manager=components_manager) - pipeline.load_components(torch_dtype=torch_dtype) - return pipeline - - def get_dummy_inputs(self, device, seed=0): - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device=device).manual_seed(seed) - inputs = { - "prompt": "A painting of a squirrel eating a burger", - "generator": generator, - "num_inference_steps": 2, - "output_type": "np", - } - return inputs - def _test_stable_diffusion_xl_euler(self, expected_image_shape, expected_slice, expected_max_diff=1e-2): - device = "cpu" # ensure determinism for the device-dependent torch.Generator sd_pipe = self.get_pipeline() - sd_pipe = sd_pipe.to(device) + sd_pipe = sd_pipe.to(torch_device) sd_pipe.set_progress_bar_config(disable=None) - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() image = sd_pipe(**inputs, output="images") image_slice = image[0, -3:, -3:, -1] assert image.shape == expected_image_shape - - assert np.abs(image_slice.flatten() - expected_slice).max() < expected_max_diff, ( - "Image Slice does not match expected slice" - ) + max_diff = torch.abs(image_slice.flatten() - expected_slice).max() + assert max_diff < expected_max_diff, f"Image slice does not match expected slice. Max Difference: {max_diff}" -class SDXLModularIPAdapterTests: +class SDXLModularIPAdapterTesterMixin: """ This mixin is designed to test IP Adapter. """ @@ -127,7 +89,7 @@ class SDXLModularIPAdapterTests: if "image" in parameters and "strength" in parameters: inputs["num_inference_steps"] = 4 - inputs["output_type"] = "np" + inputs["output_type"] = "pt" return inputs def test_ip_adapter(self, expected_max_diff: float = 1e-4, expected_pipe_slice=None): @@ -152,7 +114,7 @@ class SDXLModularIPAdapterTests: cross_attention_dim = pipe.unet.config.get("cross_attention_dim") # forward pass without ip adapter - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device)) + inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs()) if expected_pipe_slice is None: output_without_adapter = pipe(**inputs, output="images") else: @@ -163,7 +125,7 @@ class SDXLModularIPAdapterTests: pipe.unet._load_ip_adapter_weights(adapter_state_dict) # forward pass with single ip adapter, but scale=0 which should have no effect - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device)) + inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs()) inputs["ip_adapter_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)] inputs["negative_ip_adapter_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)] pipe.set_ip_adapter_scale(0.0) @@ -172,7 +134,7 @@ class SDXLModularIPAdapterTests: output_without_adapter_scale = output_without_adapter_scale[0, -3:, -3:, -1].flatten() # forward pass with single ip adapter, but with scale of adapter weights - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device)) + inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs()) inputs["ip_adapter_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)] inputs["negative_ip_adapter_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)] pipe.set_ip_adapter_scale(42.0) @@ -180,8 +142,8 @@ class SDXLModularIPAdapterTests: if expected_pipe_slice is not None: output_with_adapter_scale = output_with_adapter_scale[0, -3:, -3:, -1].flatten() - max_diff_without_adapter_scale = np.abs(output_without_adapter_scale - output_without_adapter).max() - max_diff_with_adapter_scale = np.abs(output_with_adapter_scale - output_without_adapter).max() + max_diff_without_adapter_scale = torch.abs(output_without_adapter_scale - output_without_adapter).max() + max_diff_with_adapter_scale = torch.abs(output_with_adapter_scale - output_without_adapter).max() assert max_diff_without_adapter_scale < expected_max_diff, ( "Output without ip-adapter must be same as normal inference" @@ -194,7 +156,7 @@ class SDXLModularIPAdapterTests: pipe.unet._load_ip_adapter_weights([adapter_state_dict_1, adapter_state_dict_2]) # forward pass with multi ip adapter, but scale=0 which should have no effect - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device)) + inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs()) inputs["ip_adapter_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)] * 2 inputs["negative_ip_adapter_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)] * 2 pipe.set_ip_adapter_scale([0.0, 0.0]) @@ -203,7 +165,7 @@ class SDXLModularIPAdapterTests: output_without_multi_adapter_scale = output_without_multi_adapter_scale[0, -3:, -3:, -1].flatten() # forward pass with multi ip adapter, but with scale of adapter weights - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device)) + inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs()) inputs["ip_adapter_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)] * 2 inputs["negative_ip_adapter_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)] * 2 pipe.set_ip_adapter_scale([42.0, 42.0]) @@ -211,10 +173,10 @@ class SDXLModularIPAdapterTests: if expected_pipe_slice is not None: output_with_multi_adapter_scale = output_with_multi_adapter_scale[0, -3:, -3:, -1].flatten() - max_diff_without_multi_adapter_scale = np.abs( + max_diff_without_multi_adapter_scale = torch.abs( output_without_multi_adapter_scale - output_without_adapter ).max() - max_diff_with_multi_adapter_scale = np.abs(output_with_multi_adapter_scale - output_without_adapter).max() + max_diff_with_multi_adapter_scale = torch.abs(output_with_multi_adapter_scale - output_without_adapter).max() assert max_diff_without_multi_adapter_scale < expected_max_diff, ( "Output without multi-ip-adapter must be same as normal inference" ) @@ -223,7 +185,7 @@ class SDXLModularIPAdapterTests: ) -class SDXLModularControlNetTests: +class SDXLModularControlNetTesterMixin: """ This mixin is designed to test ControlNet. """ @@ -262,24 +224,26 @@ class SDXLModularControlNetTests: pipe.set_progress_bar_config(disable=None) # forward pass without controlnet - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() output_without_controlnet = pipe(**inputs, output="images") output_without_controlnet = output_without_controlnet[0, -3:, -3:, -1].flatten() # forward pass with single controlnet, but scale=0 which should have no effect - inputs = self._modify_inputs_for_controlnet_test(self.get_dummy_inputs(torch_device)) + inputs = self._modify_inputs_for_controlnet_test(self.get_dummy_inputs()) inputs["controlnet_conditioning_scale"] = 0.0 output_without_controlnet_scale = pipe(**inputs, output="images") output_without_controlnet_scale = output_without_controlnet_scale[0, -3:, -3:, -1].flatten() # forward pass with single controlnet, but with scale of adapter weights - inputs = self._modify_inputs_for_controlnet_test(self.get_dummy_inputs(torch_device)) + inputs = self._modify_inputs_for_controlnet_test(self.get_dummy_inputs()) inputs["controlnet_conditioning_scale"] = 42.0 output_with_controlnet_scale = pipe(**inputs, output="images") output_with_controlnet_scale = output_with_controlnet_scale[0, -3:, -3:, -1].flatten() - max_diff_without_controlnet_scale = np.abs(output_without_controlnet_scale - output_without_controlnet).max() - max_diff_with_controlnet_scale = np.abs(output_with_controlnet_scale - output_without_controlnet).max() + max_diff_without_controlnet_scale = torch.abs( + output_without_controlnet_scale - output_without_controlnet + ).max() + max_diff_with_controlnet_scale = torch.abs(output_with_controlnet_scale - output_without_controlnet).max() assert max_diff_without_controlnet_scale < expected_max_diff, ( "Output without controlnet must be same as normal inference" @@ -295,21 +259,21 @@ class SDXLModularControlNetTests: guider = ClassifierFreeGuidance(guidance_scale=1.0) pipe.update_components(guider=guider) - inputs = self._modify_inputs_for_controlnet_test(self.get_dummy_inputs(torch_device)) + inputs = self._modify_inputs_for_controlnet_test(self.get_dummy_inputs()) out_no_cfg = pipe(**inputs, output="images") # forward pass with CFG applied guider = ClassifierFreeGuidance(guidance_scale=7.5) pipe.update_components(guider=guider) - inputs = self._modify_inputs_for_controlnet_test(self.get_dummy_inputs(torch_device)) + inputs = self._modify_inputs_for_controlnet_test(self.get_dummy_inputs()) out_cfg = pipe(**inputs, output="images") assert out_cfg.shape == out_no_cfg.shape - max_diff = np.abs(out_cfg - out_no_cfg).max() + max_diff = torch.abs(out_cfg - out_no_cfg).max() assert max_diff > 1e-2, "Output with CFG must be different from normal inference" -class SDXLModularGuiderTests: +class SDXLModularGuiderTesterMixin: def test_guider_cfg(self): pipe = self.get_pipeline() pipe = pipe.to(torch_device) @@ -319,13 +283,13 @@ class SDXLModularGuiderTests: guider = ClassifierFreeGuidance(guidance_scale=1.0) pipe.update_components(guider=guider) - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() out_no_cfg = pipe(**inputs, output="images") # forward pass with CFG applied guider = ClassifierFreeGuidance(guidance_scale=7.5) pipe.update_components(guider=guider) - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() out_cfg = pipe(**inputs, output="images") assert out_cfg.shape == out_no_cfg.shape @@ -333,30 +297,57 @@ class SDXLModularGuiderTests: assert max_diff > 1e-2, "Output with CFG must be different from normal inference" -class SDXLModularPipelineFastTests( - SDXLModularTests, - SDXLModularIPAdapterTests, - SDXLModularControlNetTests, - SDXLModularGuiderTests, +class TestSDXLModularPipelineFast( + SDXLModularTesterMixin, + SDXLModularIPAdapterTesterMixin, + SDXLModularControlNetTesterMixin, + SDXLModularGuiderTesterMixin, ModularPipelineTesterMixin, - unittest.TestCase, ): """Test cases for Stable Diffusion XL modular pipeline fast tests.""" + pipeline_class = StableDiffusionXLModularPipeline + pipeline_blocks_class = StableDiffusionXLAutoBlocks + repo = "hf-internal-testing/tiny-sdxl-modular" + params = frozenset( + [ + "prompt", + "height", + "width", + "negative_prompt", + "cross_attention_kwargs", + ] + ) + batch_params = frozenset(["prompt", "negative_prompt"]) + expected_image_output_shape = (1, 3, 64, 64) + + def get_dummy_inputs(self, seed=0): + generator = self.get_generator(seed) + inputs = { + "prompt": "A painting of a squirrel eating a burger", + "generator": generator, + "num_inference_steps": 2, + "output_type": "pt", + } + return inputs + def test_stable_diffusion_xl_euler(self): self._test_stable_diffusion_xl_euler( - expected_image_shape=(1, 64, 64, 3), - expected_slice=[ - 0.5966781, - 0.62939394, - 0.48465094, - 0.51573336, - 0.57593524, - 0.47035995, - 0.53410417, - 0.51436996, - 0.47313565, - ], + expected_image_shape=self.expected_image_output_shape, + expected_slice=torch.tensor( + [ + 0.5966781, + 0.62939394, + 0.48465094, + 0.51573336, + 0.57593524, + 0.47035995, + 0.53410417, + 0.51436996, + 0.47313565, + ], + device=torch_device, + ), expected_max_diff=1e-2, ) @@ -364,39 +355,65 @@ class SDXLModularPipelineFastTests( super().test_inference_batch_single_identical(expected_max_diff=3e-3) -class SDXLImg2ImgModularPipelineFastTests( - SDXLModularTests, - SDXLModularIPAdapterTests, - SDXLModularControlNetTests, - SDXLModularGuiderTests, +class TestSDXLImg2ImgModularPipelineFast( + SDXLModularTesterMixin, + SDXLModularIPAdapterTesterMixin, + SDXLModularControlNetTesterMixin, + SDXLModularGuiderTesterMixin, ModularPipelineTesterMixin, - unittest.TestCase, ): """Test cases for Stable Diffusion XL image-to-image modular pipeline fast tests.""" - def get_dummy_inputs(self, device, seed=0): - inputs = super().get_dummy_inputs(device, seed) - image = floats_tensor((1, 3, 64, 64), rng=random.Random(seed)).to(device) - image = image / 2 + 0.5 - inputs["image"] = image - inputs["strength"] = 0.8 + pipeline_class = StableDiffusionXLModularPipeline + pipeline_blocks_class = StableDiffusionXLAutoBlocks + repo = "hf-internal-testing/tiny-sdxl-modular" + params = frozenset( + [ + "prompt", + "height", + "width", + "negative_prompt", + "cross_attention_kwargs", + "image", + ] + ) + batch_params = frozenset(["prompt", "negative_prompt", "image"]) + expected_image_output_shape = (1, 3, 64, 64) + + def get_dummy_inputs(self, seed=0): + generator = self.get_generator(seed) + inputs = { + "prompt": "A painting of a squirrel eating a burger", + "generator": generator, + "num_inference_steps": 4, + "output_type": "pt", + } + image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(torch_device) + image = image.cpu().permute(0, 2, 3, 1)[0] + init_image = Image.fromarray(np.uint8(image)).convert("RGB").resize((64, 64)) + + inputs["image"] = init_image + inputs["strength"] = 0.5 return inputs def test_stable_diffusion_xl_euler(self): self._test_stable_diffusion_xl_euler( - expected_image_shape=(1, 64, 64, 3), - expected_slice=[ - 0.56943184, - 0.4702148, - 0.48048905, - 0.6235963, - 0.551138, - 0.49629188, - 0.60031277, - 0.5688907, - 0.43996853, - ], + expected_image_shape=self.expected_image_output_shape, + expected_slice=torch.tensor( + [ + 0.56943184, + 0.4702148, + 0.48048905, + 0.6235963, + 0.551138, + 0.49629188, + 0.60031277, + 0.5688907, + 0.43996853, + ], + device=torch_device, + ), expected_max_diff=1e-2, ) @@ -405,20 +422,43 @@ class SDXLImg2ImgModularPipelineFastTests( class SDXLInpaintingModularPipelineFastTests( - SDXLModularTests, - SDXLModularIPAdapterTests, - SDXLModularControlNetTests, - SDXLModularGuiderTests, + SDXLModularTesterMixin, + SDXLModularIPAdapterTesterMixin, + SDXLModularControlNetTesterMixin, + SDXLModularGuiderTesterMixin, ModularPipelineTesterMixin, - unittest.TestCase, ): """Test cases for Stable Diffusion XL inpainting modular pipeline fast tests.""" + pipeline_class = StableDiffusionXLModularPipeline + pipeline_blocks_class = StableDiffusionXLAutoBlocks + repo = "hf-internal-testing/tiny-sdxl-modular" + params = frozenset( + [ + "prompt", + "height", + "width", + "negative_prompt", + "cross_attention_kwargs", + "image", + "mask_image", + ] + ) + batch_params = frozenset(["prompt", "negative_prompt", "image", "mask_image"]) + expected_image_output_shape = (1, 3, 64, 64) + def get_dummy_inputs(self, device, seed=0): - inputs = super().get_dummy_inputs(device, seed) + generator = self.get_generator(seed) + inputs = { + "prompt": "A painting of a squirrel eating a burger", + "generator": generator, + "num_inference_steps": 4, + "output_type": "pt", + } image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device) image = image.cpu().permute(0, 2, 3, 1)[0] init_image = Image.fromarray(np.uint8(image)).convert("RGB").resize((64, 64)) + # create mask image[8:, 8:, :] = 255 mask_image = Image.fromarray(np.uint8(image)).convert("L").resize((64, 64)) @@ -431,18 +471,21 @@ class SDXLInpaintingModularPipelineFastTests( def test_stable_diffusion_xl_euler(self): self._test_stable_diffusion_xl_euler( - expected_image_shape=(1, 64, 64, 3), - expected_slice=[ - 0.40872607, - 0.38842705, - 0.34893104, - 0.47837183, - 0.43792963, - 0.5332134, - 0.3716843, - 0.47274873, - 0.45000193, - ], + expected_image_shape=self.expected_image_output_shape, + expected_slice=torch.tensor( + [ + 0.40872607, + 0.38842705, + 0.34893104, + 0.47837183, + 0.43792963, + 0.5332134, + 0.3716843, + 0.47274873, + 0.45000193, + ], + device=torch_device, + ), expected_max_diff=1e-2, ) diff --git a/tests/modular_pipelines/test_modular_pipelines_common.py b/tests/modular_pipelines/test_modular_pipelines_common.py index d309fcf353..1325e5c1de 100644 --- a/tests/modular_pipelines/test_modular_pipelines_common.py +++ b/tests/modular_pipelines/test_modular_pipelines_common.py @@ -1,9 +1,7 @@ import gc import tempfile -import unittest from typing import Callable, Union -import numpy as np import torch import diffusers @@ -19,17 +17,9 @@ from ..testing_utils import ( ) -def to_np(tensor): - if isinstance(tensor, torch.Tensor): - tensor = tensor.detach().cpu().numpy() - - return tensor - - @require_torch class ModularPipelineTesterMixin: """ - This mixin is designed to be used with unittest.TestCase classes. It provides a set of common tests for each modular pipeline, including: - test_pipeline_call_signature: check if the pipeline's __call__ method has all required parameters @@ -57,9 +47,8 @@ class ModularPipelineTesterMixin: ] ) - def get_generator(self, seed): - device = torch_device if torch_device != "mps" else "cpu" - generator = torch.Generator(device).manual_seed(seed) + def get_generator(self, seed=0): + generator = torch.Generator("cpu").manual_seed(seed) return generator @property @@ -82,13 +71,7 @@ class ModularPipelineTesterMixin: "See existing pipeline tests for reference." ) - def get_pipeline(self): - raise NotImplementedError( - "You need to implement `get_pipeline(self)` in the child test class. " - "See existing pipeline tests for reference." - ) - - def get_dummy_inputs(self, device, seed=0): + def get_dummy_inputs(self, seed=0): raise NotImplementedError( "You need to implement `get_dummy_inputs(self, device, seed)` in the child test class. " "See existing pipeline tests for reference." @@ -123,20 +106,23 @@ class ModularPipelineTesterMixin: "See existing pipeline tests for reference." ) - def setUp(self): + def setup_method(self): # clean up the VRAM before each test - super().setUp() torch.compiler.reset() gc.collect() backend_empty_cache(torch_device) - def tearDown(self): + def teardown_method(self): # clean up the VRAM after each test in case of CUDA runtime errors - super().tearDown() torch.compiler.reset() gc.collect() backend_empty_cache(torch_device) + def get_pipeline(self, components_manager=None, torch_dtype=torch.float32): + pipeline = self.pipeline_blocks_class().init_pipeline(self.repo, components_manager=components_manager) + pipeline.load_components(torch_dtype=torch_dtype) + return pipeline + def test_pipeline_call_signature(self): pipe = self.get_pipeline() input_parameters = pipe.blocks.input_names @@ -156,7 +142,7 @@ class ModularPipelineTesterMixin: pipe.to(torch_device) pipe.set_progress_bar_config(disable=None) - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["generator"] = self.get_generator(0) logger = logging.get_logger(pipe.__module__) @@ -196,7 +182,7 @@ class ModularPipelineTesterMixin: pipe = self.get_pipeline() pipe.to(torch_device) pipe.set_progress_bar_config(disable=None) - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() # Reset generator in case it is has been used in self.get_dummy_inputs inputs["generator"] = self.get_generator(0) @@ -226,10 +212,9 @@ class ModularPipelineTesterMixin: assert output_batch.shape[0] == batch_size - max_diff = np.abs(to_np(output_batch[0]) - to_np(output[0])).max() + max_diff = torch.abs(output_batch[0] - output[0]).max() assert max_diff < expected_max_diff, "Batch inference results different from single inference results" - @unittest.skipIf(torch_device not in ["cuda", "xpu"], reason="float16 requires CUDA or XPU") @require_accelerator def test_float16_inference(self, expected_max_diff=5e-2): pipe = self.get_pipeline() @@ -240,13 +225,13 @@ class ModularPipelineTesterMixin: pipe_fp16.to(torch_device, torch.float16) pipe_fp16.set_progress_bar_config(disable=None) - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() # Reset generator in case it is used inside dummy inputs if "generator" in inputs: inputs["generator"] = self.get_generator(0) output = pipe(**inputs, output="images") - fp16_inputs = self.get_dummy_inputs(torch_device) + fp16_inputs = self.get_dummy_inputs() # Reset generator in case it is used inside dummy inputs if "generator" in fp16_inputs: fp16_inputs["generator"] = self.get_generator(0) @@ -283,8 +268,8 @@ class ModularPipelineTesterMixin: pipe.set_progress_bar_config(disable=None) pipe.to("cpu") - output = pipe(**self.get_dummy_inputs("cpu"), output="images") - assert np.isnan(to_np(output)).sum() == 0, "CPU Inference returns NaN" + output = pipe(**self.get_dummy_inputs(), output="images") + assert torch.isnan(output).sum() == 0, "CPU Inference returns NaN" @require_accelerator def test_inference_is_not_nan(self): @@ -292,8 +277,8 @@ class ModularPipelineTesterMixin: pipe.set_progress_bar_config(disable=None) pipe.to(torch_device) - output = pipe(**self.get_dummy_inputs(torch_device), output="images") - assert np.isnan(to_np(output)).sum() == 0, "Accelerator Inference returns NaN" + output = pipe(**self.get_dummy_inputs(), output="images") + assert torch.isnan(output).sum() == 0, "Accelerator Inference returns NaN" def test_num_images_per_prompt(self): pipe = self.get_pipeline() @@ -309,7 +294,7 @@ class ModularPipelineTesterMixin: for batch_size in batch_sizes: for num_images_per_prompt in num_images_per_prompts: - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() for key in inputs.keys(): if key in self.batch_params: @@ -329,12 +314,12 @@ class ModularPipelineTesterMixin: image_slices = [] for pipe in [base_pipe, offload_pipe]: - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() image = pipe(**inputs, output="images") image_slices.append(image[0, -3:, -3:, -1].flatten()) - assert np.abs(image_slices[0] - image_slices[1]).max() < 1e-3 + assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3 def test_save_from_pretrained(self): pipes = [] @@ -351,9 +336,9 @@ class ModularPipelineTesterMixin: image_slices = [] for pipe in pipes: - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() image = pipe(**inputs, output="images") image_slices.append(image[0, -3:, -3:, -1].flatten()) - assert np.abs(image_slices[0] - image_slices[1]).max() < 1e-3 + assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3