From 7c451d5d99cc204c9bb2eb99c3c25388c0e8bfe5 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Thu, 16 Jul 2026 14:46:57 +0530 Subject: [PATCH] deprecate torch_dtype and prefer dtype following transformers. --- src/diffusers/loaders/ip_adapter.py | 17 +++- .../modular_pipelines/modular_pipeline.py | 6 +- .../modular_pipeline_utils.py | 14 ++- .../pipelines/pipeline_loading_utils.py | 22 ++--- src/diffusers/pipelines/pipeline_utils.py | 46 ++++++---- src/diffusers/utils/__init__.py | 2 +- src/diffusers/utils/deprecation_utils.py | 2 + .../anima/test_modular_pipeline_anima.py | 4 +- .../cosmos/test_modular_pipeline_cosmos3.py | 6 +- .../flux/test_modular_pipeline_flux.py | 8 +- ...est_modular_pipeline_stable_diffusion_3.py | 16 ++-- ...st_modular_pipeline_stable_diffusion_xl.py | 2 +- .../test_modular_pipelines_common.py | 32 +++---- .../test_modular_pipelines_custom_blocks.py | 28 ++++-- tests/pipelines/bria/test_pipeline_bria.py | 14 +-- tests/pipelines/cosmos/test_cosmos.py | 10 +-- .../cosmos/test_cosmos2_5_predict.py | 16 ++-- .../cosmos/test_cosmos2_5_transfer.py | 16 ++-- .../cosmos/test_cosmos2_text2image.py | 10 +-- .../cosmos/test_cosmos2_video2world.py | 10 +-- .../cosmos/test_cosmos_video2world.py | 10 +-- tests/pipelines/prx/test_pipeline_prx.py | 2 +- tests/pipelines/test_pipelines.py | 87 +++++++++++++++++-- tests/pipelines/test_pipelines_auto.py | 18 ++-- tests/pipelines/test_pipelines_common.py | 12 +-- 25 files changed, 261 insertions(+), 149 deletions(-) diff --git a/src/diffusers/loaders/ip_adapter.py b/src/diffusers/loaders/ip_adapter.py index cc305630b77a..5f8d3f48c997 100644 --- a/src/diffusers/loaders/ip_adapter.py +++ b/src/diffusers/loaders/ip_adapter.py @@ -29,6 +29,7 @@ is_accelerate_available, is_torch_version, is_transformers_available, + is_transformers_version, logging, ) from .unet_loader_utils import _maybe_expand_lora_scales @@ -204,13 +205,19 @@ def load_ip_adapter( else: image_encoder_subfolder = Path(image_encoder_folder).as_posix() + # transformers renamed `torch_dtype` to `dtype` in 4.56.0. + dtype_kwarg = ( + {"dtype": self.dtype} + if is_transformers_version(">=", "4.56.0") + else {"torch_dtype": self.dtype} + ) image_encoder = CLIPVisionModelWithProjection.from_pretrained( pretrained_model_name_or_path_or_dict, subfolder=image_encoder_subfolder, low_cpu_mem_usage=low_cpu_mem_usage, cache_dir=cache_dir, local_files_only=local_files_only, - torch_dtype=self.dtype, + **dtype_kwarg, ).to(self.device) self.register_modules(image_encoder=image_encoder) else: @@ -1043,11 +1050,17 @@ def load_ip_adapter( "cache_dir": cache_dir, "local_files_only": local_files_only, } + # transformers renamed `torch_dtype` to `dtype` in 4.56.0. + dtype_kwarg = ( + {"dtype": self.dtype} + if is_transformers_version(">=", "4.56.0") + else {"torch_dtype": self.dtype} + ) self.register_modules( feature_extractor=SiglipImageProcessor.from_pretrained(image_encoder_subfolder, **kwargs), image_encoder=SiglipVisionModel.from_pretrained( - image_encoder_subfolder, torch_dtype=self.dtype, **kwargs + image_encoder_subfolder, **dtype_kwarg, **kwargs ).to(self.device), ) else: diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index d43825860d8e..b4eef4c621b0 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -2347,8 +2347,8 @@ def load_components(self, names: list[str] | str | None = None, **kwargs): default_creation_method == "from_pretrained". If provided as a list or string, will load only the specified components. **kwargs: additional kwargs to be passed to `from_pretrained()`.Can be: - - a single value to be applied to all components to be loaded, e.g. torch_dtype=torch.bfloat16 - - a dict, e.g. torch_dtype={"unet": torch.bfloat16, "default": torch.float32} + - a single value to be applied to all components to be loaded, e.g. dtype=torch.bfloat16 + - a dict, e.g. dtype={"unet": torch.bfloat16, "default": torch.float32} - if potentially override ComponentSpec if passed a different loading field in kwargs, e.g. `pretrained_model_name_or_path`, `variant`, `revision`, etc. - if potentially override ComponentSpec if passed a different loading field in kwargs, e.g. @@ -2638,7 +2638,7 @@ def module_is_offloaded(module): " is not recommended to move them to `cpu` as running them will fail. Please make" " sure to use an accelerator to run the pipeline in inference, due to the lack of" " support for`float16` operations on this device in PyTorch. Please, remove the" - " `torch_dtype=torch.float16` argument, or use another device for inference." + " `dtype=torch.float16` argument, or use another device for inference." ) return self diff --git a/src/diffusers/modular_pipelines/modular_pipeline_utils.py b/src/diffusers/modular_pipelines/modular_pipeline_utils.py index 7aa61c55ef1f..b6e82e67c015 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline_utils.py +++ b/src/diffusers/modular_pipelines/modular_pipeline_utils.py @@ -26,7 +26,7 @@ from ..configuration_utils import ConfigMixin, FrozenDict from ..loaders.single_file_utils import _is_single_file_path_or_url -from ..utils import DIFFUSERS_LOAD_ID_FIELDS, is_torch_available, logging +from ..utils import _TORCH_DTYPE_DEPRECATION_MESSAGE, DIFFUSERS_LOAD_ID_FIELDS, deprecate, is_torch_available, logging from ..utils.import_utils import _is_package_available @@ -293,6 +293,15 @@ def create(self, config: FrozenDict | dict[str, Any] | None = None, **kwargs) -> # YiYi TODO: add guard for type of model, if it is supported by from_pretrained def load(self, **kwargs) -> Any: """Load component using from_pretrained.""" + torch_dtype = kwargs.pop("torch_dtype", None) + if torch_dtype is not None: + if kwargs.get("dtype") is not None: + raise ValueError( + "You have passed both `dtype` and `torch_dtype`. Please only pass `dtype`, `torch_dtype` is deprecated." + ) + deprecate("torch_dtype", "1.0.0", _TORCH_DTYPE_DEPRECATION_MESSAGE) + kwargs["dtype"] = torch_dtype + # select loading fields from kwargs passed from user: e.g. pretrained_model_name_or_path, subfolder, variant, revision, note the list could change passed_loading_kwargs = {key: kwargs.pop(key) for key in self.loading_fields() if key in kwargs} # merge loading field value in the spec with user passed values to create load_kwargs @@ -311,12 +320,11 @@ def load(self, **kwargs) -> Any: from diffusers import AutoModel - # `dtype`/`torch_dtype` is not an accepted parameter for tokenizers and processors. + # `dtype` is not an accepted parameter for tokenizers and processors. # As a result, it gets stored in `init_kwargs`, which are written to the config # during save. This causes JSON serialization to fail when saving the component. if self.type_hint is not None and not issubclass(self.type_hint, (torch.nn.Module, AutoModel)): kwargs.pop("dtype", None) - kwargs.pop("torch_dtype", None) if self.type_hint is None: try: diff --git a/src/diffusers/pipelines/pipeline_loading_utils.py b/src/diffusers/pipelines/pipeline_loading_utils.py index 90cbffc5b69d..7856d5714254 100644 --- a/src/diffusers/pipelines/pipeline_loading_utils.py +++ b/src/diffusers/pipelines/pipeline_loading_utils.py @@ -563,7 +563,7 @@ def _load_empty_model( pipelines: Any, is_pipeline_module: bool, name: str, - torch_dtype: str | torch.dtype, + dtype: str | torch.dtype, cached_folder: str | os.PathLike, **kwargs, ): @@ -636,7 +636,7 @@ def _load_empty_model( model = class_obj(config) if model is not None: - model = model.to(dtype=torch_dtype) + model = model.to(dtype=dtype) return model @@ -677,7 +677,7 @@ def _get_final_device_map(device_map, pipeline_class, passed_class_obj, init_dic # To avoid circular import problem. from diffusers import pipelines - torch_dtype = kwargs.get("torch_dtype", torch.float32) + dtype = kwargs.get("dtype", torch.float32) # Load each module in the pipeline on a meta device so that we can derive the device map. init_empty_modules = {} @@ -708,9 +708,7 @@ def _get_final_device_map(device_map, pipeline_class, passed_class_obj, init_dic else: sub_model_dtype = ( - torch_dtype.get(name, torch_dtype.get("default", torch.float32)) - if isinstance(torch_dtype, dict) - else torch_dtype + dtype.get(name, dtype.get("default", torch.float32)) if isinstance(dtype, dict) else dtype ) loaded_sub_model = _load_empty_model( library_name=library_name, @@ -720,7 +718,7 @@ def _get_final_device_map(device_map, pipeline_class, passed_class_obj, init_dic is_pipeline_module=is_pipeline_module, pipeline_class=pipeline_class, name=name, - torch_dtype=sub_model_dtype, + dtype=sub_model_dtype, cached_folder=kwargs.get("cached_folder", None), force_download=kwargs.get("force_download", None), proxies=kwargs.get("proxies", None), @@ -738,9 +736,7 @@ def _get_final_device_map(device_map, pipeline_class, passed_class_obj, init_dic module_sizes = { module_name: compute_module_sizes( module, - dtype=torch_dtype.get(module_name, torch_dtype.get("default", torch.float32)) - if isinstance(torch_dtype, dict) - else torch_dtype, + dtype=dtype.get(module_name, dtype.get("default", torch.float32)) if isinstance(dtype, dict) else dtype, )[""] for module_name, module in init_empty_modules.items() if isinstance(module, torch.nn.Module) @@ -776,7 +772,7 @@ def load_sub_model( pipelines: Any, is_pipeline_module: bool, pipeline_class: Any, - torch_dtype: torch.dtype, + dtype: torch.dtype, provider: Any, sess_options: Any, device_map: dict[str, torch.device] | str | None, @@ -856,9 +852,9 @@ def load_sub_model( # For transformers models >= 4.56.0, use 'dtype' instead of 'torch_dtype' to avoid deprecation warnings if issubclass(class_obj, torch.nn.Module): if is_transformers_model and transformers_version >= version.parse("4.56.0"): - loading_kwargs["dtype"] = torch_dtype + loading_kwargs["dtype"] = dtype else: - loading_kwargs["torch_dtype"] = torch_dtype + loading_kwargs["torch_dtype"] = dtype if issubclass(class_obj, diffusers_module.OnnxRuntimeModel): loading_kwargs["provider"] = provider loading_kwargs["sess_options"] = sess_options diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index d737b44129ea..2fe9b4f4fd17 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -59,6 +59,7 @@ from ..quantizers.bitsandbytes.utils import _check_bnb_status from ..schedulers.scheduling_utils import SCHEDULER_CONFIG_NAME from ..utils import ( + _TORCH_DTYPE_DEPRECATION_MESSAGE, CONFIG_NAME, DEPRECATED_REVISION_ARGS, BaseOutput, @@ -592,7 +593,7 @@ def module_is_offloaded(module): " is not recommended to move them to `cpu` as running them will fail. Please make" " sure to use an accelerator to run the pipeline in inference, due to the lack of" " support for`float16` operations on this device in PyTorch. Please, remove the" - " `torch_dtype=torch.float16` argument, or use another device for inference." + " `dtype=torch.float16` argument, or use another device for inference." ) return self @@ -784,7 +785,13 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike, **kwa from_flax = kwargs.pop("from_flax", False) torch_dtype = kwargs.pop("torch_dtype", None) dtype = kwargs.pop("dtype", None) - torch_dtype = dtype if dtype is not None else torch_dtype + if torch_dtype is not None: + if dtype is not None: + raise ValueError( + "You have passed both `dtype` and `torch_dtype`. Please only pass `dtype`, `torch_dtype` is deprecated." + ) + deprecate("torch_dtype", "1.0.0", _TORCH_DTYPE_DEPRECATION_MESSAGE) + dtype = torch_dtype custom_pipeline = kwargs.pop("custom_pipeline", None) custom_revision = kwargs.pop("custom_revision", None) provider = kwargs.pop("provider", None) @@ -805,11 +812,9 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike, **kwa disable_mmap = kwargs.pop("disable_mmap", False) trust_remote_code = kwargs.pop("trust_remote_code", False) - if torch_dtype is not None and not isinstance(torch_dtype, dict) and not isinstance(torch_dtype, torch.dtype): - torch_dtype = torch.float32 - logger.warning( - f"Passed `torch_dtype` {torch_dtype} is not a `torch.dtype`. Defaulting to `torch.float32`." - ) + if dtype is not None and not isinstance(dtype, dict) and not isinstance(dtype, torch.dtype): + logger.warning(f"Passed `dtype` {dtype} is not a `torch.dtype`. Defaulting to `torch.float32`.") + dtype = torch.float32 if low_cpu_mem_usage and not is_accelerate_available(): low_cpu_mem_usage = False @@ -1019,7 +1024,7 @@ def load_module(name, value): init_dict=init_dict, library=library, max_memory=max_memory, - torch_dtype=torch_dtype, + dtype=dtype, cached_folder=cached_folder, force_download=force_download, proxies=proxies, @@ -1067,9 +1072,7 @@ def load_module(name, value): else: # load sub model sub_model_dtype = ( - torch_dtype.get(name, torch_dtype.get("default", torch.float32)) - if isinstance(torch_dtype, dict) - else torch_dtype + dtype.get(name, dtype.get("default", torch.float32)) if isinstance(dtype, dict) else dtype ) loaded_sub_model = load_sub_model( library_name=library_name, @@ -1078,7 +1081,7 @@ def load_module(name, value): pipelines=pipelines, is_pipeline_module=is_pipeline_module, pipeline_class=pipeline_class, - torch_dtype=sub_model_dtype, + dtype=sub_model_dtype, provider=provider, sess_options=sess_options, device_map=current_device_map, @@ -1466,7 +1469,7 @@ def enable_group_offload( >>> from diffusers import DiffusionPipeline >>> import torch - >>> pipe = DiffusionPipeline.from_pretrained("Qwen/Qwen-Image", torch_dtype=torch.bfloat16) + >>> pipe = DiffusionPipeline.from_pretrained("Qwen/Qwen-Image", dtype=torch.bfloat16) >>> pipe.enable_group_offload( ... onload_device=torch.device("cuda"), @@ -2055,7 +2058,7 @@ def enable_xformers_memory_efficient_attention(self, attention_op: Callable | No >>> from diffusers import DiffusionPipeline >>> from xformers.ops import MemoryEfficientAttentionFlashAttentionOp - >>> pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16) + >>> pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", dtype=torch.float16) >>> pipe = pipe.to("cuda") >>> pipe.enable_xformers_memory_efficient_attention(attention_op=MemoryEfficientAttentionFlashAttentionOp) >>> # Workaround for not accepting attention shape using VAE for Flash Attention @@ -2114,7 +2117,7 @@ def enable_attention_slicing(self, slice_size: str | int = "auto"): >>> pipe = StableDiffusionPipeline.from_pretrained( ... "stable-diffusion-v1-5/stable-diffusion-v1-5", - ... torch_dtype=torch.float16, + ... dtype=torch.float16, ... use_safetensors=True, ... ) @@ -2168,7 +2171,15 @@ def from_pipe(cls, pipeline, **kwargs): original_config = dict(pipeline.config) torch_dtype = kwargs.pop("torch_dtype", None) dtype = kwargs.pop("dtype", None) - torch_dtype = dtype if dtype is not None else (torch_dtype if torch_dtype is not None else torch.float32) + if torch_dtype is not None: + if dtype is not None: + raise ValueError( + "You have passed both `dtype` and `torch_dtype`. Please only pass `dtype`, `torch_dtype` is deprecated." + ) + deprecate("torch_dtype", "1.0.0", _TORCH_DTYPE_DEPRECATION_MESSAGE) + dtype = torch_dtype + if dtype is None: + dtype = torch.float32 trust_remote_code = kwargs.pop("trust_remote_code", False) # derive the pipeline class to instantiate @@ -2266,8 +2277,7 @@ def from_pipe(cls, pipeline, **kwargs): new_pipeline.register_to_config(_name_or_path=pretrained_model_name_or_path) new_pipeline.register_to_config(**unused_original_config) - if torch_dtype is not None: - new_pipeline.to(dtype=torch_dtype) + new_pipeline.to(dtype=dtype) return new_pipeline diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py index 5cd6885e0364..a14931148d28 100644 --- a/src/diffusers/utils/__init__.py +++ b/src/diffusers/utils/__init__.py @@ -41,7 +41,7 @@ WEIGHTS_INDEX_NAME, WEIGHTS_NAME, ) -from .deprecation_utils import _maybe_remap_transformers_class, deprecate +from .deprecation_utils import _TORCH_DTYPE_DEPRECATION_MESSAGE, _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 encode_video, 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 7ec612499e9a..45041b0a2877 100644 --- a/src/diffusers/utils/deprecation_utils.py +++ b/src/diffusers/utils/deprecation_utils.py @@ -9,6 +9,8 @@ logger = logging.get_logger(__name__) +_TORCH_DTYPE_DEPRECATION_MESSAGE = "Please use `dtype` instead." + # 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 diff --git a/tests/modular_pipelines/anima/test_modular_pipeline_anima.py b/tests/modular_pipelines/anima/test_modular_pipeline_anima.py index 9afe9ee6fab3..253c5175951e 100644 --- a/tests/modular_pipelines/anima/test_modular_pipeline_anima.py +++ b/tests/modular_pipelines/anima/test_modular_pipeline_anima.py @@ -155,10 +155,10 @@ class TestAnimaModularPipelineFast(ModularPipelineTesterMixin, ModularGuiderTest batch_params = frozenset(["prompt", "negative_prompt"]) expected_workflow_blocks = ANIMA_TEXT2IMAGE_WORKFLOWS - def get_pipeline(self, components_manager=None, torch_dtype=torch.float32): + def get_pipeline(self, components_manager=None, dtype=torch.float32): pipe = self.pipeline_blocks_class().init_pipeline(components_manager=components_manager) pipe.update_components(**get_dummy_components()) - pipe.to(dtype=torch_dtype) + pipe.to(dtype=dtype) pipe.set_progress_bar_config(disable=None) return pipe diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py index 5297a441cebf..480f7d3be58e 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py @@ -122,8 +122,8 @@ class TestCosmos3OmniModularPipelineFast(ModularPipelineTesterMixin): output_name = "videos" expected_workflow_blocks = COSMOS3_OMNI_WORKFLOWS - def get_pipeline(self, components_manager=None, torch_dtype=torch.float32): - pipe = super().get_pipeline(components_manager, torch_dtype) + def get_pipeline(self, components_manager=None, dtype=torch.float32): + pipe = super().get_pipeline(components_manager, dtype) pipe.disable_safety_checker() return pipe @@ -161,7 +161,7 @@ def test_save_from_pretrained(self, tmp_path): base_pipe.save_pretrained(str(tmp_path)) loaded_pipe = ModularPipeline.from_pretrained(str(tmp_path)) - loaded_pipe.load_components(torch_dtype=torch.float32) + loaded_pipe.load_components(dtype=torch.float32) loaded_pipe.disable_safety_checker() loaded_pipe.to(torch_device) diff --git a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py index d82a9e2f395b..00df9865e1c3 100644 --- a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py +++ b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py @@ -98,8 +98,8 @@ class TestFluxImg2ImgModularPipelineFast(ModularPipelineTesterMixin): batch_params = frozenset(["prompt", "image"]) expected_workflow_blocks = FLUX_IMAGE2IMAGE_WORKFLOWS - def get_pipeline(self, components_manager=None, torch_dtype=torch.float32): - pipeline = super().get_pipeline(components_manager, torch_dtype) + def get_pipeline(self, components_manager=None, dtype=torch.float32): + pipeline = super().get_pipeline(components_manager, dtype) # Override `vae_scale_factor` here as currently, `image_processor` is initialized with # fixed constants instead of @@ -135,7 +135,7 @@ def test_save_from_pretrained(self, tmp_path): base_pipe.save_pretrained(str(tmp_path)) pipe = ModularPipeline.from_pretrained(tmp_path).to(torch_device) - pipe.load_components(torch_dtype=torch.float32) + pipe.load_components(dtype=torch.float32) pipe.to(torch_device) pipe.image_processor = VaeImageProcessor(vae_scale_factor=2) @@ -216,7 +216,7 @@ def test_save_from_pretrained(self, tmp_path): base_pipe.save_pretrained(str(tmp_path)) pipe = ModularPipeline.from_pretrained(tmp_path).to(torch_device) - pipe.load_components(torch_dtype=torch.float32) + pipe.load_components(dtype=torch.float32) pipe.to(torch_device) pipe.image_processor = VaeImageProcessor(vae_scale_factor=2) diff --git a/tests/modular_pipelines/stable_diffusion_3/test_modular_pipeline_stable_diffusion_3.py b/tests/modular_pipelines/stable_diffusion_3/test_modular_pipeline_stable_diffusion_3.py index d9cffcf6c36d..52965a653935 100644 --- a/tests/modular_pipelines/stable_diffusion_3/test_modular_pipeline_stable_diffusion_3.py +++ b/tests/modular_pipelines/stable_diffusion_3/test_modular_pipeline_stable_diffusion_3.py @@ -67,8 +67,8 @@ def get_dummy_inputs(self, seed=0): "output_type": "pt", } - def get_pipeline(self, components_manager=None, torch_dtype=torch.float32): - return super().get_pipeline(components_manager, torch_dtype) + def get_pipeline(self, components_manager=None, dtype=torch.float32): + return super().get_pipeline(components_manager, dtype) def test_save_from_pretrained(self, tmp_path): pipes = [] @@ -77,7 +77,7 @@ def test_save_from_pretrained(self, tmp_path): base_pipe.save_pretrained(str(tmp_path)) pipe = self.pipeline_class.from_pretrained(tmp_path).to(torch_device) - pipe.load_components(torch_dtype=torch.float32) + pipe.load_components(dtype=torch.float32) pipe.to(torch_device) pipes.append(pipe) @@ -94,7 +94,7 @@ def test_load_expected_components_from_save_pretrained(self, tmp_path): base_pipe.save_pretrained(str(tmp_path)) pipe = self.pipeline_class.from_pretrained(tmp_path) - pipe.load_components(torch_dtype=torch.float32) + pipe.load_components(dtype=torch.float32) assert set(base_pipe.components.keys()) == set(pipe.components.keys()) @@ -135,8 +135,8 @@ def test_pipeline_call_signature(self): # (guidance_scale) which are intentionally omitted from pipeline inputs. pass - def get_pipeline(self, components_manager=None, torch_dtype=torch.float32): - pipeline = super().get_pipeline(components_manager, torch_dtype) + def get_pipeline(self, components_manager=None, dtype=torch.float32): + pipeline = super().get_pipeline(components_manager, dtype) pipeline.image_processor = VaeImageProcessor(vae_scale_factor=8) return pipeline @@ -165,7 +165,7 @@ def test_save_from_pretrained(self, tmp_path): base_pipe.save_pretrained(str(tmp_path)) pipe = self.pipeline_class.from_pretrained(tmp_path).to(torch_device) - pipe.load_components(torch_dtype=torch.float32) + pipe.load_components(dtype=torch.float32) pipe.to(torch_device) pipe.image_processor = VaeImageProcessor(vae_scale_factor=8) pipes.append(pipe) @@ -183,7 +183,7 @@ def test_load_expected_components_from_save_pretrained(self, tmp_path): base_pipe.save_pretrained(str(tmp_path)) pipe = self.pipeline_class.from_pretrained(tmp_path) - pipe.load_components(torch_dtype=torch.float32) + pipe.load_components(dtype=torch.float32) assert set(base_pipe.components.keys()) == set(pipe.components.keys()) 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 7bbfac596365..d6a0523c2d11 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 @@ -106,7 +106,7 @@ def test_ip_adapter(self, expected_max_diff: float = 1e-4, expected_pipe_slice=N blocks = self.pipeline_blocks_class() _ = blocks.sub_blocks.pop("ip_adapter") pipe = blocks.init_pipeline(self.pretrained_model_name_or_path) - pipe.load_components(torch_dtype=torch.float32) + pipe.load_components(dtype=torch.float32) pipe = pipe.to(torch_device) cross_attention_dim = pipe.unet.config.get("cross_attention_dim") diff --git a/tests/modular_pipelines/test_modular_pipelines_common.py b/tests/modular_pipelines/test_modular_pipelines_common.py index 223a25e436fa..2e3495f30c0c 100644 --- a/tests/modular_pipelines/test_modular_pipelines_common.py +++ b/tests/modular_pipelines/test_modular_pipelines_common.py @@ -154,11 +154,11 @@ def teardown_method(self): gc.collect() backend_empty_cache(torch_device) - def get_pipeline(self, components_manager=None, torch_dtype=torch.float32): + def get_pipeline(self, components_manager=None, dtype=torch.float32): pipeline = self.pipeline_blocks_class().init_pipeline( self.pretrained_model_name_or_path, components_manager=components_manager ) - pipeline.load_components(torch_dtype=torch_dtype) + pipeline.load_components(dtype=dtype) pipeline.set_progress_bar_config(disable=None) return pipeline @@ -369,7 +369,7 @@ def test_save_from_pretrained(self, tmp_path): base_pipe.save_pretrained(str(tmp_path)) pipe = ModularPipeline.from_pretrained(tmp_path).to(torch_device) - pipe.load_components(torch_dtype=torch.float32) + pipe.load_components(dtype=torch.float32) pipe.to(torch_device) pipes.append(pipe) @@ -403,7 +403,7 @@ def test_load_expected_components_from_save_pretrained(self, tmp_path): expected = _get_specified_components(save_dir) loaded_pipe = ModularPipeline.from_pretrained(save_dir) - loaded_pipe.load_components(torch_dtype=torch.float32) + loaded_pipe.load_components(dtype=torch.float32) actual = { name @@ -714,7 +714,7 @@ def test_automodel_tags_load_id(self): def test_automodel_update_components(self): pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe") - pipe.load_components(torch_dtype=torch.float32) + pipe.load_components(dtype=torch.float32) auto_model = AutoModel.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe", subfolder="unet") @@ -751,7 +751,7 @@ def test_load_components_loads_local_single_file_path(self, tmp_path): class TestLoadComponentsSkipBehavior: def test_load_components_skips_already_loaded(self): pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe") - pipe.load_components(torch_dtype=torch.float32) + pipe.load_components(dtype=torch.float32) original_unet = pipe.unet @@ -763,7 +763,7 @@ def test_load_components_skips_already_loaded(self): def test_load_components_selective_loading(self): pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe") - pipe.load_components(names="unet", torch_dtype=torch.float32) + pipe.load_components(names="unet", dtype=torch.float32) # Verify only requested component was loaded. assert hasattr(pipe, "unet") @@ -774,8 +774,8 @@ def test_load_components_selective_loading_incremental(self): """Loading a subset of components should not affect already-loaded components.""" pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe") - pipe.load_components(names="unet", torch_dtype=torch.float32) - pipe.load_components(names="text_encoder", torch_dtype=torch.float32) + pipe.load_components(names="unet", dtype=torch.float32) + pipe.load_components(names="text_encoder", dtype=torch.float32) assert hasattr(pipe, "unet") assert pipe.unet is not None @@ -791,7 +791,7 @@ def test_load_components_skips_invalid_pretrained_path(self): pretrained_model_name_or_path=None, default_creation_method="from_pretrained", ) - pipe.load_components(torch_dtype=torch.float32) + pipe.load_components(dtype=torch.float32) # Verify test_component was not loaded assert not hasattr(pipe, "test_component") or pipe.test_component is None @@ -804,7 +804,7 @@ def test_save_pretrained_updates_index_for_local_model(self, tmp_path): import json pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe") - pipe.load_components(torch_dtype=torch.float32) + pipe.load_components(dtype=torch.float32) pipe.unet._diffusers_load_id = "null" @@ -824,7 +824,7 @@ def test_save_pretrained_updates_index_for_local_model(self, tmp_path): def test_save_pretrained_roundtrip_with_local_model(self, tmp_path): """A pipeline with a custom/local model should be saveable and re-loadable with identical outputs.""" pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe") - pipe.load_components(torch_dtype=torch.float32) + pipe.load_components(dtype=torch.float32) pipe.unet._diffusers_load_id = "null" @@ -834,7 +834,7 @@ def test_save_pretrained_roundtrip_with_local_model(self, tmp_path): pipe.save_pretrained(save_dir) loaded_pipe = ModularPipeline.from_pretrained(save_dir) - loaded_pipe.load_components(torch_dtype=torch.float32) + loaded_pipe.load_components(dtype=torch.float32) assert loaded_pipe.unet is not None assert loaded_pipe.unet.__class__.__name__ == pipe.unet.__class__.__name__ @@ -852,7 +852,7 @@ def test_save_pretrained_updates_index_for_model_with_no_load_id(self, tmp_path) from diffusers import UNet2DConditionModel pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe") - pipe.load_components(torch_dtype=torch.float32) + pipe.load_components(dtype=torch.float32) unet = UNet2DConditionModel.from_pretrained( "hf-internal-testing/tiny-stable-diffusion-xl-pipe", subfolder="unet" @@ -879,7 +879,7 @@ def test_save_pretrained_overwrite_modular_index(self, tmp_path): import json pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe") - pipe.load_components(torch_dtype=torch.float32) + pipe.load_components(dtype=torch.float32) save_dir = str(tmp_path / "my-pipeline") pipe.save_pretrained(save_dir, overwrite_modular_index=True) @@ -897,7 +897,7 @@ def test_save_pretrained_overwrite_modular_index(self, tmp_path): assert spec["subfolder"] == component_name loaded_pipe = ModularPipeline.from_pretrained(save_dir) - loaded_pipe.load_components(torch_dtype=torch.float32) + loaded_pipe.load_components(dtype=torch.float32) assert loaded_pipe.unet is not None assert loaded_pipe.vae is not None diff --git a/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py b/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py index 433d1e872a7d..2f104c53fb20 100644 --- a/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py +++ b/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py @@ -19,6 +19,7 @@ from typing import List import numpy as np +import pytest import torch from diffusers import FluxTransformer2DModel @@ -379,8 +380,8 @@ def __call__(self, components, state: PipelineState) -> PipelineState: loaded_pipe.update_components(custom_model=custom_model) assert getattr(loaded_pipe, "custom_model", None) is not None - def test_automodel_type_hint_preserves_torch_dtype(self, tmp_path): - """Regression test for #13271: torch_dtype was incorrectly removed when type_hint is AutoModel.""" + def test_automodel_type_hint_preserves_dtype(self, tmp_path): + """Regression test for #13271: dtype was incorrectly removed when type_hint is AutoModel.""" from diffusers import AutoModel model_dir = str(tmp_path / "model") @@ -412,10 +413,27 @@ def __call__(self, components, state: PipelineState) -> PipelineState: block = DtypeTestBlock() pipe = block.init_pipeline() - pipe.load_components(torch_dtype=torch.float16, trust_remote_code=True) + pipe.load_components(dtype=torch.float16, trust_remote_code=True) assert pipe.model.dtype == torch.float16 + def test_component_spec_load_torch_dtype_is_deprecated(self, tmp_path): + from diffusers import AutoModel + + model_dir = str(tmp_path / "model") + os.makedirs(model_dir) + _create_tiny_model_dir(model_dir) + + spec = ComponentSpec("model", AutoModel, pretrained_model_name_or_path=model_dir) + + with pytest.warns(FutureWarning, match="torch_dtype"): + component = spec.load(torch_dtype=torch.float16, trust_remote_code=True) + # The deprecated argument is still honoured until it is removed. + assert component.dtype == torch.float16 + + with pytest.raises(ValueError, match="passed both `dtype` and `torch_dtype`"): + spec.load(dtype=torch.float16, torch_dtype=torch.float16, trust_remote_code=True) + @require_torch_accelerator def test_automodel_type_hint_preserves_device(self, tmp_path): """Test that ComponentSpec with AutoModel type_hint correctly passes device_map.""" @@ -594,7 +612,7 @@ def test_loading_from_hub(self): pipe.load_components( trust_remote_code=True, device_map="cuda", - torch_dtype={"default": torch.bfloat16, "vae": torch.float16}, + dtype={"default": torch.bfloat16, "vae": torch.float16}, ) assert len(pipe.components) == 7 assert sorted(pipe.components) == sorted( @@ -607,7 +625,7 @@ def test_forward(self): pipe.load_components( trust_remote_code=True, device_map="cuda", - torch_dtype={"default": torch.bfloat16, "vae": torch.float16}, + dtype={"default": torch.bfloat16, "vae": torch.float16}, ) num_frames_per_block = 2 diff --git a/tests/pipelines/bria/test_pipeline_bria.py b/tests/pipelines/bria/test_pipeline_bria.py index dac9c428cfc9..96a5f557d5fe 100644 --- a/tests/pipelines/bria/test_pipeline_bria.py +++ b/tests/pipelines/bria/test_pipeline_bria.py @@ -170,7 +170,7 @@ def test_save_load_float16(self, expected_max_diff=1e-2): with tempfile.TemporaryDirectory() as tmpdir: pipe.save_pretrained(tmpdir) - pipe_loaded = self.pipeline_class.from_pretrained(tmpdir, torch_dtype=torch.float16) + pipe_loaded = self.pipeline_class.from_pretrained(tmpdir, dtype=torch.float16) for component in pipe_loaded.components.values(): if hasattr(component, "set_default_attn_processor"): component.set_default_attn_processor() @@ -215,14 +215,14 @@ def test_to_dtype(self): model_dtypes = [component.dtype for component in components.values() if hasattr(component, "dtype")] self.assertTrue([dtype == torch.float32 for dtype in model_dtypes] == [True, True, True]) - def test_torch_dtype_dict(self): + def test_dtype_dict(self): components = self.get_dummy_components() pipe = self.pipeline_class(**components) with tempfile.TemporaryDirectory() as tmpdirname: pipe.save_pretrained(tmpdirname) - torch_dtype_dict = {"transformer": torch.bfloat16, "default": torch.float16} - loaded_pipe = self.pipeline_class.from_pretrained(tmpdirname, torch_dtype=torch_dtype_dict) + dtype_dict = {"transformer": torch.bfloat16, "default": torch.float16} + loaded_pipe = self.pipeline_class.from_pretrained(tmpdirname, dtype=dtype_dict) self.assertEqual(loaded_pipe.transformer.dtype, torch.bfloat16) self.assertEqual(loaded_pipe.text_encoder.dtype, torch.float16) @@ -230,8 +230,8 @@ def test_torch_dtype_dict(self): with tempfile.TemporaryDirectory() as tmpdirname: pipe.save_pretrained(tmpdirname) - torch_dtype_dict = {"default": torch.float16} - loaded_pipe = self.pipeline_class.from_pretrained(tmpdirname, torch_dtype=torch_dtype_dict) + dtype_dict = {"default": torch.float16} + loaded_pipe = self.pipeline_class.from_pretrained(tmpdirname, dtype=dtype_dict) self.assertEqual(loaded_pipe.transformer.dtype, torch.float16) self.assertEqual(loaded_pipe.text_encoder.dtype, torch.float16) @@ -272,7 +272,7 @@ def get_inputs(self, device, seed=0): def test_bria_inference_bf16(self): pipe = self.pipeline_class.from_pretrained( - self.repo_id, torch_dtype=torch.bfloat16, text_encoder=None, tokenizer=None + self.repo_id, dtype=torch.bfloat16, text_encoder=None, tokenizer=None ) pipe.to(torch_device) diff --git a/tests/pipelines/cosmos/test_cosmos.py b/tests/pipelines/cosmos/test_cosmos.py index 80b1fdc37c6d..85e9716b23fa 100644 --- a/tests/pipelines/cosmos/test_cosmos.py +++ b/tests/pipelines/cosmos/test_cosmos.py @@ -321,7 +321,7 @@ def test_serialization_with_variants(self): is_folder = os.path.isdir(folder_path) and subfolder in config assert is_folder and any(p.split(".")[1].startswith(variant) for p in os.listdir(folder_path)) - def test_torch_dtype_dict(self): + def test_dtype_dict(self): components = self.get_dummy_components() if not components: self.skipTest("No dummy components defined.") @@ -332,16 +332,16 @@ def test_torch_dtype_dict(self): with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname: pipe.save_pretrained(tmpdirname, safe_serialization=False) - torch_dtype_dict = {specified_key: torch.bfloat16, "default": torch.float16} + dtype_dict = {specified_key: torch.bfloat16, "default": torch.float16} loaded_pipe = self.pipeline_class.from_pretrained( - tmpdirname, safety_checker=DummyCosmosSafetyChecker(), torch_dtype=torch_dtype_dict + tmpdirname, safety_checker=DummyCosmosSafetyChecker(), dtype=dtype_dict ) for name, component in loaded_pipe.components.items(): if name == "safety_checker": continue if isinstance(component, torch.nn.Module) and hasattr(component, "dtype"): - expected_dtype = torch_dtype_dict.get(name, torch_dtype_dict.get("default", torch.float32)) + expected_dtype = dtype_dict.get(name, dtype_dict.get("default", torch.float32)) self.assertEqual( component.dtype, expected_dtype, @@ -349,7 +349,7 @@ def test_torch_dtype_dict(self): ) def test_dtype_alias(self): - # `dtype` is an alias for `torch_dtype` in `from_pretrained`. + # `torch_dtype` is deprecated in favor of `dtype` in `from_pretrained`. components = self.get_dummy_components() pipe = self.pipeline_class(**components) diff --git a/tests/pipelines/cosmos/test_cosmos2_5_predict.py b/tests/pipelines/cosmos/test_cosmos2_5_predict.py index a7475f7b2d8f..32954d74683f 100644 --- a/tests/pipelines/cosmos/test_cosmos2_5_predict.py +++ b/tests/pipelines/cosmos/test_cosmos2_5_predict.py @@ -44,9 +44,9 @@ def from_pretrained(*args, **kwargs): if "safety_checker" not in kwargs or kwargs["safety_checker"] is None: safety_checker = DummyCosmosSafetyChecker() device_map = kwargs.get("device_map", "cpu") - torch_dtype = kwargs.get("torch_dtype") - if device_map is not None or torch_dtype is not None: - safety_checker = safety_checker.to(device_map, dtype=torch_dtype) + dtype = kwargs.get("dtype") + if device_map is not None or dtype is not None: + safety_checker = safety_checker.to(device_map, dtype=dtype) kwargs["safety_checker"] = safety_checker return Cosmos2_5_PredictBasePipeline.from_pretrained(*args, **kwargs) @@ -300,7 +300,7 @@ def test_serialization_with_variants(self): is_folder = os.path.isdir(folder_path) and subfolder in config assert is_folder and any(p.split(".")[1].startswith(variant) for p in os.listdir(folder_path)) - def test_torch_dtype_dict(self): + def test_dtype_dict(self): components = self.get_dummy_components() if not components: self.skipTest("No dummy components defined.") @@ -311,16 +311,16 @@ def test_torch_dtype_dict(self): with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname: pipe.save_pretrained(tmpdirname, safe_serialization=False) - torch_dtype_dict = {specified_key: torch.bfloat16, "default": torch.float16} + dtype_dict = {specified_key: torch.bfloat16, "default": torch.float16} loaded_pipe = self.pipeline_class.from_pretrained( - tmpdirname, safety_checker=DummyCosmosSafetyChecker(), torch_dtype=torch_dtype_dict + tmpdirname, safety_checker=DummyCosmosSafetyChecker(), dtype=dtype_dict ) for name, component in loaded_pipe.components.items(): if name == "safety_checker": continue if isinstance(component, torch.nn.Module) and hasattr(component, "dtype"): - expected_dtype = torch_dtype_dict.get(name, torch_dtype_dict.get("default", torch.float32)) + expected_dtype = dtype_dict.get(name, dtype_dict.get("default", torch.float32)) self.assertEqual( component.dtype, expected_dtype, @@ -328,7 +328,7 @@ def test_torch_dtype_dict(self): ) def test_dtype_alias(self): - # `dtype` is an alias for `torch_dtype` in `from_pretrained`. + # `torch_dtype` is deprecated in favor of `dtype` in `from_pretrained`. components = self.get_dummy_components() pipe = self.pipeline_class(**components) diff --git a/tests/pipelines/cosmos/test_cosmos2_5_transfer.py b/tests/pipelines/cosmos/test_cosmos2_5_transfer.py index c59ae946733e..f95caabb4243 100644 --- a/tests/pipelines/cosmos/test_cosmos2_5_transfer.py +++ b/tests/pipelines/cosmos/test_cosmos2_5_transfer.py @@ -45,9 +45,9 @@ def from_pretrained(*args, **kwargs): if "safety_checker" not in kwargs or kwargs["safety_checker"] is None: safety_checker = DummyCosmosSafetyChecker() device_map = kwargs.get("device_map", "cpu") - torch_dtype = kwargs.get("torch_dtype") - if device_map is not None or torch_dtype is not None: - safety_checker = safety_checker.to(device_map, dtype=torch_dtype) + dtype = kwargs.get("dtype") + if device_map is not None or dtype is not None: + safety_checker = safety_checker.to(device_map, dtype=dtype) kwargs["safety_checker"] = safety_checker return Cosmos2_5_TransferPipeline.from_pretrained(*args, **kwargs) @@ -397,7 +397,7 @@ def test_serialization_with_variants(self): is_folder = os.path.isdir(folder_path) and subfolder in config assert is_folder and any(p.split(".")[1].startswith(variant) for p in os.listdir(folder_path)) - def test_torch_dtype_dict(self): + def test_dtype_dict(self): components = self.get_dummy_components() if not components: self.skipTest("No dummy components defined.") @@ -408,9 +408,9 @@ def test_torch_dtype_dict(self): with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname: pipe.save_pretrained(tmpdirname, safe_serialization=False) - torch_dtype_dict = {specified_key: torch.bfloat16, "default": torch.float16} + dtype_dict = {specified_key: torch.bfloat16, "default": torch.float16} loaded_pipe = self.pipeline_class.from_pretrained( - tmpdirname, safety_checker=DummyCosmosSafetyChecker(), torch_dtype=torch_dtype_dict + tmpdirname, safety_checker=DummyCosmosSafetyChecker(), dtype=dtype_dict ) for name, component in loaded_pipe.components.items(): @@ -418,7 +418,7 @@ def test_torch_dtype_dict(self): if name == "safety_checker": continue if isinstance(component, torch.nn.Module) and hasattr(component, "dtype"): - expected_dtype = torch_dtype_dict.get(name, torch_dtype_dict.get("default", torch.float32)) + expected_dtype = dtype_dict.get(name, dtype_dict.get("default", torch.float32)) self.assertEqual( component.dtype, expected_dtype, @@ -426,7 +426,7 @@ def test_torch_dtype_dict(self): ) def test_dtype_alias(self): - # `dtype` is an alias for `torch_dtype` in `from_pretrained`. + # `torch_dtype` is deprecated in favor of `dtype` in `from_pretrained`. components = self.get_dummy_components() pipe = self.pipeline_class(**components) diff --git a/tests/pipelines/cosmos/test_cosmos2_text2image.py b/tests/pipelines/cosmos/test_cosmos2_text2image.py index eb624365bc42..232d5faf1d02 100644 --- a/tests/pipelines/cosmos/test_cosmos2_text2image.py +++ b/tests/pipelines/cosmos/test_cosmos2_text2image.py @@ -305,7 +305,7 @@ def test_serialization_with_variants(self): is_folder = os.path.isdir(folder_path) and subfolder in config assert is_folder and any(p.split(".")[1].startswith(variant) for p in os.listdir(folder_path)) - def test_torch_dtype_dict(self): + def test_dtype_dict(self): components = self.get_dummy_components() if not components: self.skipTest("No dummy components defined.") @@ -316,16 +316,16 @@ def test_torch_dtype_dict(self): with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname: pipe.save_pretrained(tmpdirname, safe_serialization=False) - torch_dtype_dict = {specified_key: torch.bfloat16, "default": torch.float16} + dtype_dict = {specified_key: torch.bfloat16, "default": torch.float16} loaded_pipe = self.pipeline_class.from_pretrained( - tmpdirname, safety_checker=DummyCosmosSafetyChecker(), torch_dtype=torch_dtype_dict + tmpdirname, safety_checker=DummyCosmosSafetyChecker(), dtype=dtype_dict ) for name, component in loaded_pipe.components.items(): if name == "safety_checker": continue if isinstance(component, torch.nn.Module) and hasattr(component, "dtype"): - expected_dtype = torch_dtype_dict.get(name, torch_dtype_dict.get("default", torch.float32)) + expected_dtype = dtype_dict.get(name, dtype_dict.get("default", torch.float32)) self.assertEqual( component.dtype, expected_dtype, @@ -333,7 +333,7 @@ def test_torch_dtype_dict(self): ) def test_dtype_alias(self): - # `dtype` is an alias for `torch_dtype` in `from_pretrained`. + # `torch_dtype` is deprecated in favor of `dtype` in `from_pretrained`. components = self.get_dummy_components() pipe = self.pipeline_class(**components) diff --git a/tests/pipelines/cosmos/test_cosmos2_video2world.py b/tests/pipelines/cosmos/test_cosmos2_video2world.py index 6baf872c8a91..a951877db2b1 100644 --- a/tests/pipelines/cosmos/test_cosmos2_video2world.py +++ b/tests/pipelines/cosmos/test_cosmos2_video2world.py @@ -319,7 +319,7 @@ def test_serialization_with_variants(self): is_folder = os.path.isdir(folder_path) and subfolder in config assert is_folder and any(p.split(".")[1].startswith(variant) for p in os.listdir(folder_path)) - def test_torch_dtype_dict(self): + def test_dtype_dict(self): components = self.get_dummy_components() if not components: self.skipTest("No dummy components defined.") @@ -330,16 +330,16 @@ def test_torch_dtype_dict(self): with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname: pipe.save_pretrained(tmpdirname, safe_serialization=False) - torch_dtype_dict = {specified_key: torch.bfloat16, "default": torch.float16} + dtype_dict = {specified_key: torch.bfloat16, "default": torch.float16} loaded_pipe = self.pipeline_class.from_pretrained( - tmpdirname, safety_checker=DummyCosmosSafetyChecker(), torch_dtype=torch_dtype_dict + tmpdirname, safety_checker=DummyCosmosSafetyChecker(), dtype=dtype_dict ) for name, component in loaded_pipe.components.items(): if name == "safety_checker": continue if isinstance(component, torch.nn.Module) and hasattr(component, "dtype"): - expected_dtype = torch_dtype_dict.get(name, torch_dtype_dict.get("default", torch.float32)) + expected_dtype = dtype_dict.get(name, dtype_dict.get("default", torch.float32)) self.assertEqual( component.dtype, expected_dtype, @@ -347,7 +347,7 @@ def test_torch_dtype_dict(self): ) def test_dtype_alias(self): - # `dtype` is an alias for `torch_dtype` in `from_pretrained`. + # `torch_dtype` is deprecated in favor of `dtype` in `from_pretrained`. components = self.get_dummy_components() pipe = self.pipeline_class(**components) diff --git a/tests/pipelines/cosmos/test_cosmos_video2world.py b/tests/pipelines/cosmos/test_cosmos_video2world.py index e49cf8a41939..7e4928ccf0df 100644 --- a/tests/pipelines/cosmos/test_cosmos_video2world.py +++ b/tests/pipelines/cosmos/test_cosmos_video2world.py @@ -334,7 +334,7 @@ def test_serialization_with_variants(self): is_folder = os.path.isdir(folder_path) and subfolder in config assert is_folder and any(p.split(".")[1].startswith(variant) for p in os.listdir(folder_path)) - def test_torch_dtype_dict(self): + def test_dtype_dict(self): components = self.get_dummy_components() if not components: self.skipTest("No dummy components defined.") @@ -345,16 +345,16 @@ def test_torch_dtype_dict(self): with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname: pipe.save_pretrained(tmpdirname, safe_serialization=False) - torch_dtype_dict = {specified_key: torch.bfloat16, "default": torch.float16} + dtype_dict = {specified_key: torch.bfloat16, "default": torch.float16} loaded_pipe = self.pipeline_class.from_pretrained( - tmpdirname, safety_checker=DummyCosmosSafetyChecker(), torch_dtype=torch_dtype_dict + tmpdirname, safety_checker=DummyCosmosSafetyChecker(), dtype=dtype_dict ) for name, component in loaded_pipe.components.items(): if name == "safety_checker": continue if isinstance(component, torch.nn.Module) and hasattr(component, "dtype"): - expected_dtype = torch_dtype_dict.get(name, torch_dtype_dict.get("default", torch.float32)) + expected_dtype = dtype_dict.get(name, dtype_dict.get("default", torch.float32)) self.assertEqual( component.dtype, expected_dtype, @@ -362,7 +362,7 @@ def test_torch_dtype_dict(self): ) def test_dtype_alias(self): - # `dtype` is an alias for `torch_dtype` in `from_pretrained`. + # `torch_dtype` is deprecated in favor of `dtype` in `from_pretrained`. components = self.get_dummy_components() pipe = self.pipeline_class(**components) diff --git a/tests/pipelines/prx/test_pipeline_prx.py b/tests/pipelines/prx/test_pipeline_prx.py index a7f05a61d52e..104f7b0b7553 100644 --- a/tests/pipelines/prx/test_pipeline_prx.py +++ b/tests/pipelines/prx/test_pipeline_prx.py @@ -274,7 +274,7 @@ def test_save_load_optional_components(self): pass @unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.") - def test_torch_dtype_dict(self): + def test_dtype_dict(self): pass @unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.") diff --git a/tests/pipelines/test_pipelines.py b/tests/pipelines/test_pipelines.py index 8aa874e7aa6d..e89cefdc0cb0 100644 --- a/tests/pipelines/test_pipelines.py +++ b/tests/pipelines/test_pipelines.py @@ -1342,7 +1342,7 @@ def test_download_from_git(self): custom_pipeline="clip_guided_stable_diffusion", clip_model=clip_model, feature_extractor=feature_extractor, - torch_dtype=torch.float16, + dtype=torch.float16, ) pipeline.enable_attention_slicing() pipeline = pipeline.to(torch_device) @@ -1938,7 +1938,6 @@ def test_pipe_to(self): sd4 = sd.to(device=device_type) sd5 = sd.to(torch_device=device_type) sd6 = sd.to(device_type, dtype=torch.float32) - sd7 = sd.to(device_type, torch_dtype=torch.float32) assert sd1.device.type == device_type assert sd2.device.type == device_type @@ -1946,33 +1945,27 @@ def test_pipe_to(self): assert sd4.device.type == device_type assert sd5.device.type == device_type assert sd6.device.type == device_type - assert sd7.device.type == device_type sd1 = sd.to(torch.float16) sd2 = sd.to(None, torch.float16) sd3 = sd.to(dtype=torch.float16) sd4 = sd.to(dtype=torch.float16) sd5 = sd.to(None, dtype=torch.float16) - sd6 = sd.to(None, torch_dtype=torch.float16) assert sd1.dtype == torch.float16 assert sd2.dtype == torch.float16 assert sd3.dtype == torch.float16 assert sd4.dtype == torch.float16 assert sd5.dtype == torch.float16 - assert sd6.dtype == torch.float16 sd1 = sd.to(device=device_type, dtype=torch.float16) - sd2 = sd.to(torch_device=device_type, torch_dtype=torch.float16) - sd3 = sd.to(device_type, torch.float16) + sd2 = sd.to(device_type, torch.float16) assert sd1.dtype == torch.float16 assert sd2.dtype == torch.float16 - assert sd3.dtype == torch.float16 assert sd1.device.type == device_type assert sd2.device.type == device_type - assert sd3.device.type == device_type def test_pipe_same_device_id_offload(self): unet = self.dummy_cond_unet() @@ -2011,6 +2004,82 @@ def test_dduf_file_is_deprecated(self): ) assert "dduf_file" in str(warning_ctx.warning) + def test_torch_dtype_is_deprecated(self): + sd = StableDiffusionPipeline( + unet=self.dummy_cond_unet(), + scheduler=PNDMScheduler(skip_prk_steps=True), + vae=self.dummy_vae, + text_encoder=self.dummy_text_encoder, + tokenizer=CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"), + safety_checker=None, + feature_extractor=self.dummy_extractor, + ) + + with tempfile.TemporaryDirectory() as tmpdirname: + sd.save_pretrained(tmpdirname) + with self.assertWarns(FutureWarning) as warning_ctx: + loaded_sd = StableDiffusionPipeline.from_pretrained(tmpdirname, torch_dtype=torch.float16) + + assert "torch_dtype" in str(warning_ctx.warning) + assert "Please use `dtype` instead." in str(warning_ctx.warning) + # The deprecated argument is still honoured until it is removed. + assert loaded_sd.unet.dtype == torch.float16 + + def test_dtype_does_not_warn(self): + sd = StableDiffusionPipeline( + unet=self.dummy_cond_unet(), + scheduler=PNDMScheduler(skip_prk_steps=True), + vae=self.dummy_vae, + text_encoder=self.dummy_text_encoder, + tokenizer=CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"), + safety_checker=None, + feature_extractor=self.dummy_extractor, + ) + + with tempfile.TemporaryDirectory() as tmpdirname: + sd.save_pretrained(tmpdirname) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + loaded_sd = StableDiffusionPipeline.from_pretrained(tmpdirname, dtype=torch.float16) + + assert not [w for w in caught if issubclass(w.category, FutureWarning) and "torch_dtype" in str(w.message)] + assert loaded_sd.unet.dtype == torch.float16 + + def test_torch_dtype_and_dtype_raises(self): + sd = StableDiffusionPipeline( + unet=self.dummy_cond_unet(), + scheduler=PNDMScheduler(skip_prk_steps=True), + vae=self.dummy_vae, + text_encoder=self.dummy_text_encoder, + tokenizer=CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"), + safety_checker=None, + feature_extractor=self.dummy_extractor, + ) + + with tempfile.TemporaryDirectory() as tmpdirname: + sd.save_pretrained(tmpdirname) + with self.assertRaises(ValueError) as error_context: + _ = StableDiffusionPipeline.from_pretrained(tmpdirname, dtype=torch.float16, torch_dtype=torch.float16) + + assert "passed both `dtype` and `torch_dtype`" in str(error_context.exception) + + def test_from_pipe_torch_dtype_is_deprecated(self): + sd = StableDiffusionPipeline( + unet=self.dummy_cond_unet(), + scheduler=PNDMScheduler(skip_prk_steps=True), + vae=self.dummy_vae, + text_encoder=self.dummy_text_encoder, + tokenizer=CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"), + safety_checker=None, + feature_extractor=self.dummy_extractor, + ) + + with self.assertWarns(FutureWarning) as warning_ctx: + img2img = StableDiffusionImg2ImgPipeline.from_pipe(sd, torch_dtype=torch.float16) + + assert "torch_dtype" in str(warning_ctx.warning) + assert img2img.unet.dtype == torch.float16 + @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") diff --git a/tests/pipelines/test_pipelines_auto.py b/tests/pipelines/test_pipelines_auto.py index 3d6c607b1a32..0502967c6a1b 100644 --- a/tests/pipelines/test_pipelines_auto.py +++ b/tests/pipelines/test_pipelines_auto.py @@ -467,9 +467,7 @@ class AutoPipelineIntegrationTest(unittest.TestCase): def test_pipe_auto(self): for model_name, model_repo in PRETRAINED_MODEL_REPO_MAPPING.items(): # test txt2img - pipe_txt2img = AutoPipelineForText2Image.from_pretrained( - model_repo, variant="fp16", torch_dtype=torch.float16 - ) + pipe_txt2img = AutoPipelineForText2Image.from_pretrained(model_repo, variant="fp16", dtype=torch.float16) self.assertIsInstance(pipe_txt2img, AUTO_TEXT2IMAGE_PIPELINES_MAPPING[model_name]) pipe_to = AutoPipelineForText2Image.from_pipe(pipe_txt2img) @@ -487,9 +485,7 @@ def test_pipe_auto(self): # test img2img - pipe_img2img = AutoPipelineForImage2Image.from_pretrained( - model_repo, variant="fp16", torch_dtype=torch.float16 - ) + pipe_img2img = AutoPipelineForImage2Image.from_pretrained(model_repo, variant="fp16", dtype=torch.float16) self.assertIsInstance(pipe_img2img, AUTO_IMAGE2IMAGE_PIPELINES_MAPPING[model_name]) pipe_to = AutoPipelineForText2Image.from_pipe(pipe_img2img) @@ -509,7 +505,7 @@ def test_pipe_auto(self): if "kandinsky" not in model_name: pipe_inpaint = AutoPipelineForInpainting.from_pretrained( - model_repo, variant="fp16", torch_dtype=torch.float16 + model_repo, variant="fp16", dtype=torch.float16 ) self.assertIsInstance(pipe_inpaint, AUTO_INPAINT_PIPELINES_MAPPING[model_name]) @@ -534,7 +530,7 @@ def test_from_pipe_consistent(self): # test from_pretrained for pipe_from_class in auto_pipes: - pipe_from = pipe_from_class.from_pretrained(model_repo, variant="fp16", torch_dtype=torch.float16) + pipe_from = pipe_from_class.from_pretrained(model_repo, variant="fp16", dtype=torch.float16) pipe_from_config = dict(pipe_from.config) for pipe_to_class in auto_pipes: @@ -552,17 +548,17 @@ def test_controlnet(self): controlnet = ControlNetModel.from_pretrained(controlnet_repo, torch_dtype=torch.float16) pipe_txt2img = AutoPipelineForText2Image.from_pretrained( - model_repo, controlnet=controlnet, torch_dtype=torch.float16 + model_repo, controlnet=controlnet, dtype=torch.float16 ) self.assertIsInstance(pipe_txt2img, AUTO_TEXT2IMAGE_PIPELINES_MAPPING["stable-diffusion-controlnet"]) pipe_img2img = AutoPipelineForImage2Image.from_pretrained( - model_repo, controlnet=controlnet, torch_dtype=torch.float16 + model_repo, controlnet=controlnet, dtype=torch.float16 ) self.assertIsInstance(pipe_img2img, AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["stable-diffusion-controlnet"]) pipe_inpaint = AutoPipelineForInpainting.from_pretrained( - model_repo, controlnet=controlnet, torch_dtype=torch.float16 + model_repo, controlnet=controlnet, dtype=torch.float16 ) self.assertIsInstance(pipe_inpaint, AUTO_INPAINT_PIPELINES_MAPPING["stable-diffusion-controlnet"]) diff --git a/tests/pipelines/test_pipelines_common.py b/tests/pipelines/test_pipelines_common.py index 4c8a9c6ddc47..20e5e46fb08b 100644 --- a/tests/pipelines/test_pipelines_common.py +++ b/tests/pipelines/test_pipelines_common.py @@ -1470,7 +1470,7 @@ def test_save_load_float16(self, expected_max_diff=1e-2): with tempfile.TemporaryDirectory() as tmpdir: pipe.save_pretrained(tmpdir) - pipe_loaded = self.pipeline_class.from_pretrained(tmpdir, torch_dtype=torch.float16) + pipe_loaded = self.pipeline_class.from_pretrained(tmpdir, dtype=torch.float16) for component in pipe_loaded.components.values(): if hasattr(component, "set_default_attn_processor"): component.set_default_attn_processor() @@ -2347,7 +2347,7 @@ def run_forward(pipe): self.assertTrue(np.allclose(output_without_group_offloading, output_with_group_offloading1, atol=1e-4)) self.assertTrue(np.allclose(output_without_group_offloading, output_with_group_offloading2, atol=1e-4)) - def test_torch_dtype_dict(self): + def test_dtype_dict(self): components = self.get_dummy_components() if not components: self.skipTest("No dummy components defined.") @@ -2357,12 +2357,12 @@ def test_torch_dtype_dict(self): with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname: pipe.save_pretrained(tmpdirname, safe_serialization=False) - torch_dtype_dict = {specified_key: torch.bfloat16, "default": torch.float16} - loaded_pipe = self.pipeline_class.from_pretrained(tmpdirname, torch_dtype=torch_dtype_dict) + dtype_dict = {specified_key: torch.bfloat16, "default": torch.float16} + loaded_pipe = self.pipeline_class.from_pretrained(tmpdirname, dtype=dtype_dict) for name, component in loaded_pipe.components.items(): if isinstance(component, torch.nn.Module) and hasattr(component, "dtype"): - expected_dtype = torch_dtype_dict.get(name, torch_dtype_dict.get("default", torch.float32)) + expected_dtype = dtype_dict.get(name, dtype_dict.get("default", torch.float32)) self.assertEqual( component.dtype, expected_dtype, @@ -2370,7 +2370,7 @@ def test_torch_dtype_dict(self): ) def test_dtype_alias(self): - # `dtype` is an alias for `torch_dtype` in `from_pretrained`. + # `torch_dtype` is deprecated in favor of `dtype` in `from_pretrained`. components = self.get_dummy_components() pipe = self.pipeline_class(**components)