# coding=utf-8 # Copyright 2022 The HuggingFace Inc. team. # Copyright (c) 2022, NVIDIA CORPORATION. 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 copy import importlib import inspect import os import sys from typing import Optional, Union import torch from diffusers.pipelines import DiffusionPipeline from diffusers.pipelines.pipeline_utils import _unwrap_model from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card from diffusers.utils.torch_utils import is_compiled_module from huggingface_hub import create_repo from optimum.utils import logging from ...transformers.gaudi_configuration import GaudiConfig from ...utils import to_device_dtype logger = logging.get_logger(__name__) GAUDI_LOADABLE_CLASSES = { "diffusers": { "ModelMixin": ["save_pretrained", "from_pretrained"], "SchedulerMixin": ["save_pretrained", "from_pretrained"], "DiffusionPipeline": ["save_pretrained", "from_pretrained"], "OnnxRuntimeModel": ["save_pretrained", "from_pretrained"], }, "transformers": { "PreTrainedTokenizer": ["save_pretrained", "from_pretrained"], "PreTrainedTokenizerFast": ["save_pretrained", "from_pretrained"], "PreTrainedModel": ["save_pretrained", "from_pretrained"], "FeatureExtractionMixin": ["save_pretrained", "from_pretrained"], "ProcessorMixin": ["save_pretrained", "from_pretrained"], "ImageProcessingMixin": ["save_pretrained", "from_pretrained"], }, "optimum.habana.diffusers.schedulers": { "GaudiDDIMScheduler": ["save_pretrained", "from_pretrained"], "GaudiEulerDiscreteScheduler": ["save_pretrained", "from_pretrained"], "GaudiFlowMatchEulerDiscreteScheduler": ["save_pretrained", "from_pretrained"], "GaudiEulerAncestralDiscreteScheduler": ["save_pretrained", "from_pretrained"], }, } GAUDI_ALL_IMPORTABLE_CLASSES = {} for library in GAUDI_LOADABLE_CLASSES: GAUDI_ALL_IMPORTABLE_CLASSES.update(GAUDI_LOADABLE_CLASSES[library]) def _fetch_class_library_tuple(module): # import it here to avoid circular import from diffusers import pipelines # register the config from the original module, not the dynamo compiled one not_compiled_module = _unwrap_model(module) library = not_compiled_module.__module__.split(".")[0] if library == "optimum": library = "optimum.habana.diffusers.schedulers" # check if the module is a pipeline module module_path_items = not_compiled_module.__module__.split(".") pipeline_dir = module_path_items[-2] if len(module_path_items) > 2 else None path = not_compiled_module.__module__.split(".") is_pipeline_module = pipeline_dir in path and hasattr(pipelines, pipeline_dir) # if library is not in GAUDI_LOADABLE_CLASSES, then it is a custom module. # Or if it's a pipeline module, then the module is inside the pipeline # folder so we set the library to module name. if is_pipeline_module: library = pipeline_dir elif library not in GAUDI_LOADABLE_CLASSES: library = not_compiled_module.__module__ # retrieve class_name class_name = not_compiled_module.__class__.__name__ return (library, class_name) class GaudiDiffusionPipeline(DiffusionPipeline): """ Extends the [`DiffusionPipeline`](https://huggingface.co/docs/diffusers/api/diffusion_pipeline) class: - The pipeline is initialized on Gaudi if `use_habana=True`. - The pipeline's Gaudi configuration is saved and pushed to the hub. Args: use_habana (bool, defaults to `False`): Whether to use Gaudi (`True`) or CPU (`False`). use_hpu_graphs (bool, defaults to `False`): Whether to use HPU graphs or not. gaudi_config (Union[str, [`GaudiConfig`]], defaults to `None`): Gaudi configuration to use. Can be a string to download it from the Hub. Or a previously initialized config can be passed. bf16_full_eval (bool, defaults to `False`): Whether to use full bfloat16 evaluation instead of 32-bit. This will be faster and save memory compared to fp32/mixed precision but can harm generated images. sdp_on_bf16 (bool, defaults to `False`): Whether to allow PyTorch to use reduced precision in the SDPA math backend. """ def __init__( self, use_habana: bool = False, use_hpu_graphs: bool = False, gaudi_config: Union[str, GaudiConfig] = None, bf16_full_eval: bool = False, sdp_on_bf16: bool = False, ): DiffusionPipeline.__init__(self) if sdp_on_bf16: if hasattr(torch._C, "_set_math_sdp_allow_fp16_bf16_reduction"): torch._C._set_math_sdp_allow_fp16_bf16_reduction(True) self.use_habana = use_habana if self.use_habana: self.use_hpu_graphs = use_hpu_graphs if self.use_hpu_graphs: logger.info("Enabled HPU graphs.") else: logger.info("Enabled lazy mode because `use_hpu_graphs=False`.") self._device = torch.device("hpu") import diffusers # Patch for unconditional image generation from ..models import gaudi_unet_2d_model_forward diffusers.models.unets.unet_2d.UNet2DModel.forward = gaudi_unet_2d_model_forward if isinstance(gaudi_config, str): # Config from the Hub self.gaudi_config = GaudiConfig.from_pretrained(gaudi_config) elif isinstance(gaudi_config, GaudiConfig): # Config already initialized self.gaudi_config = copy.deepcopy(gaudi_config) else: raise ValueError( f"`gaudi_config` must be a string or a GaudiConfig object but is {type(gaudi_config)}." ) if self.gaudi_config.use_torch_autocast: if bf16_full_eval: logger.warning( "`use_torch_autocast` is True in the given Gaudi configuration but " "`torch_dtype=torch.bfloat16` was given. Disabling mixed precision and continuing in bf16 only." ) self.gaudi_config.use_torch_autocast = False # Workaround for Synapse 1.11 for full bf16 and Torch Autocast if bf16_full_eval or self.gaudi_config.use_torch_autocast: import diffusers from ..models import ( gaudi_unet_2d_condition_model_forward, ) diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.forward = ( gaudi_unet_2d_condition_model_forward ) if self.use_hpu_graphs: try: import habana_frameworks.torch as ht except ImportError as error: error.msg = f"Could not import habana_frameworks.torch. {error.msg}." raise error self.ht = ht self.hpu_stream = ht.hpu.Stream() self.cache = {} else: try: import habana_frameworks.torch.core as htcore except ImportError as error: error.msg = f"Could not import habana_frameworks.torch.core. {error.msg}." raise error self.htcore = htcore else: if use_hpu_graphs: raise ValueError( "`use_hpu_graphs` is True but `use_habana` is False, please set `use_habana=True` to use HPU" " graphs." ) if gaudi_config is not None: raise ValueError( "Got a non-None `gaudi_config` but `use_habana` is False, please set `use_habana=True` to use this" " Gaudi configuration." ) logger.info("Running on CPU.") self._device = torch.device("cpu") def register_modules(self, **kwargs): for name, module in kwargs.items(): # retrieve library if module is None or isinstance(module, (tuple, list)) and module[0] is None: register_dict = {name: (None, None)} else: library, class_name = _fetch_class_library_tuple(module) register_dict = {name: (library, class_name)} # save model index config self.register_to_config(**register_dict) # set models setattr(self, name, module) def save_pretrained( self, save_directory: Union[str, os.PathLike], safe_serialization: bool = True, variant: Optional[str] = None, push_to_hub: bool = False, **kwargs, ): """ Save the pipeline and Gaudi configurations. More information [here](https://huggingface.co/docs/diffusers/api/diffusion_pipeline#diffusers.DiffusionPipeline.save_pretrained). Arguments: save_directory (`str` or `os.PathLike`): Directory to which to save. Will be created if it doesn't exist. safe_serialization (`bool`, *optional*, defaults to `True`): Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`). variant (`str`, *optional*): If specified, weights are saved in the format pytorch_model..bin. push_to_hub (`bool`, *optional*, defaults to `False`): Whether or not to push your model to the Hugging Face model 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. """ model_index_dict = dict(self.config) model_index_dict.pop("_class_name", None) model_index_dict.pop("_diffusers_version", None) model_index_dict.pop("_module", None) model_index_dict.pop("_name_or_path", None) if push_to_hub: commit_message = kwargs.pop("commit_message", None) private = kwargs.pop("private", False) 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 expected_modules, optional_kwargs = self._get_signature_keys(self) def is_saveable_module(name, value): if name not in expected_modules: return False if name in self._optional_components and value[0] is None: return False return True model_index_dict = {k: v for k, v in model_index_dict.items() if is_saveable_module(k, v)} for pipeline_component_name in model_index_dict.keys(): sub_model = getattr(self, pipeline_component_name) model_cls = sub_model.__class__ # Dynamo wraps the original model in a private class. # I didn't find a public API to get the original class. if is_compiled_module(sub_model): sub_model = _unwrap_model(sub_model) model_cls = sub_model.__class__ save_method_name = None # search for the model's base class in GAUDI_LOADABLE_CLASSES for library_name, library_classes in GAUDI_LOADABLE_CLASSES.items(): if library_name in sys.modules: library = importlib.import_module(library_name) else: logger.info( f"{library_name} is not installed. Cannot save {pipeline_component_name} as {library_classes} from {library_name}" ) for base_class, save_load_methods in library_classes.items(): class_candidate = getattr(library, base_class, None) if class_candidate is not None and issubclass(model_cls, class_candidate): # if we found a suitable base class in GAUDI_LOADABLE_CLASSES then grab its save method save_method_name = save_load_methods[0] break if save_method_name is not None: break if save_method_name is None: logger.warn(f"self.{pipeline_component_name}={sub_model} of type {type(sub_model)} cannot be saved.") # make sure that unsaveable components are not tried to be loaded afterward self.register_to_config(**{pipeline_component_name: (None, None)}) continue save_method = getattr(sub_model, save_method_name) # Call the save method with the argument safe_serialization only if it's supported save_method_signature = inspect.signature(save_method) save_method_accept_safe = "safe_serialization" in save_method_signature.parameters save_method_accept_variant = "variant" in save_method_signature.parameters save_kwargs = {} if save_method_accept_safe: save_kwargs["safe_serialization"] = safe_serialization if save_method_accept_variant: save_kwargs["variant"] = variant save_method(os.path.join(save_directory, pipeline_component_name), **save_kwargs) # finally save the config self.save_config(save_directory) if hasattr(self, "gaudi_config"): self.gaudi_config.save_pretrained(save_directory) if push_to_hub: # Create a new empty model card and eventually tag it model_card = load_or_create_model_card(repo_id, token=token, is_pipeline=True) model_card = populate_model_card(model_card) model_card.save(os.path.join(save_directory, "README.md")) self._upload_folder( save_directory, repo_id, token=token, commit_message=commit_message, create_pr=create_pr, ) def to(self, *args, **kwargs): """ Intercept to() method and disable gpu-hpu migration before sending to diffusers """ kwargs["hpu_migration"] = False return super().to( *args, **kwargs, ) @classmethod def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs): """ More information [here](https://huggingface.co/docs/diffusers/api/diffusion_pipeline#diffusers.DiffusionPipeline.from_pretrained). """ # Set the correct log level depending on the node # Already done in super().init() but we have to do it again # because we use optimum.utils.logging here and not # diffusers.utils.logging log_level = kwargs.pop("log_level", logging.INFO) logging.set_verbosity(log_level) logging.enable_default_handler() logging.enable_explicit_format() # Import diffusers.pipelines.pipeline_utils to override the values of LOADABLE_CLASSES and ALL_IMPORTABLE_CLASSES import diffusers.pipelines.pipeline_utils diffusers.pipelines.pipeline_utils.LOADABLE_CLASSES = GAUDI_LOADABLE_CLASSES diffusers.pipelines.pipeline_utils.ALL_IMPORTABLE_CLASSES = GAUDI_ALL_IMPORTABLE_CLASSES # Define a new kwarg here to know in the __init__ whether to use full bf16 precision or not bf16_full_eval = kwargs.get("torch_dtype", None) == torch.bfloat16 kwargs["bf16_full_eval"] = bf16_full_eval # Need to load custom ops lists before instantiating htcore if kwargs.get("gaudi_config", None) is not None: if isinstance(kwargs["gaudi_config"], str): gaudi_config = GaudiConfig.from_pretrained(kwargs["gaudi_config"]) else: gaudi_config = kwargs["gaudi_config"] gaudi_config.declare_autocast_bf16_fp32_ops() kwargs["gaudi_config"] = gaudi_config # Import htcore here to support model quantization import habana_frameworks.torch.core as htcore # noqa: F401 # Normally we just need to return super().from_pretrained. However this is a # workaround for Transformers 4.49.0 issue (sub_model torch_dtype option ignored). # Note this issue is already fixed in 4.50.0dev working branch.. model = super().from_pretrained( pretrained_model_name_or_path, **kwargs, ) if bf16_full_eval: # Get the component names component_names = [name for name in model.__dict__ if not name.startswith("_")] # Iterate through the component names and fix dtype for name in component_names: component = getattr(model, name, None) if component is not None and hasattr(component, "dtype"): component.to(torch.bfloat16) return model @classmethod def save_lora_weights( cls, save_directory: Union[str, os.PathLike], **kwargs, ): # Move all lora layers state dicts from HPU to CPU before saving for key in list(kwargs.keys()): if key.endswith("_lora_layers") and kwargs[key] is not None: kwargs[key] = to_device_dtype(kwargs[key], target_device=torch.device("cpu")) # Call diffusers' base class handler return super().save_lora_weights( save_directory, **kwargs, ) """ def save_lora_weights( cls, save_directory: Union[str, os.PathLike], unet_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, text_encoder_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, text_encoder_2_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, ): # Move the state dict from HPU to CPU before saving if unet_lora_layers: unet_lora_layers = to_device_dtype(unet_lora_layers, target_device=torch.device("cpu")) if text_encoder_lora_layers: text_encoder_lora_layers = to_device_dtype(text_encoder_lora_layers, target_device=torch.device("cpu")) if text_encoder_2_lora_layers: text_encoder_2_lora_layers = to_device_dtype(text_encoder_2_lora_layers, target_device=torch.device("cpu")) # text_encoder_2_lora_layers is only supported by some diffuser pipelines signature = inspect.signature(super().save_lora_weights) if "text_encoder_2_lora_layers" in signature.parameters: return super().save_lora_weights( save_directory, unet_lora_layers, text_encoder_lora_layers, text_encoder_2_lora_layers, is_main_process, weight_name, save_function, safe_serialization, ) else: return super().save_lora_weights( save_directory, unet_lora_layers, text_encoder_lora_layers, is_main_process, weight_name, save_function, safe_serialization, ) """