diff --git a/src/diffusers/loaders/single_file_utils.py b/src/diffusers/loaders/single_file_utils.py index ef6c41e3ce97..2b67f91dc205 100644 --- a/src/diffusers/loaders/single_file_utils.py +++ b/src/diffusers/loaders/single_file_utils.py @@ -387,6 +387,14 @@ def is_valid_url(url): return False +def _validate_single_file_path(pretrained_model_name_or_path): + if os.path.isfile(pretrained_model_name_or_path): + return True + + repo_id, weight_name = _extract_repo_id_and_weights_name(pretrained_model_name_or_path) + return bool(repo_id and weight_name) + + def _extract_repo_id_and_weights_name(pretrained_model_name_or_path): if not is_valid_url(pretrained_model_name_or_path): raise ValueError("Invalid `pretrained_model_name_or_path` provided. Please set it to a valid URL.") @@ -398,7 +406,6 @@ def _extract_repo_id_and_weights_name(pretrained_model_name_or_path): pretrained_model_name_or_path = pretrained_model_name_or_path.replace(prefix, "") match = re.match(pattern, pretrained_model_name_or_path) if not match: - logger.warning("Unable to identify the repo_id and weights_name from the provided URL.") return repo_id, weights_name repo_id = f"{match.group(1)}/{match.group(2)}" diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 78226a49b122..27b8ab0e42d1 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -366,7 +366,7 @@ def init_pipeline( collection: Optional[str] = None, ) -> "ModularPipeline": """ - create a ModularPipeline, optionally accept modular_repo to load from hub. + create a ModularPipeline, optionally accept pretrained_model_name_or_path to load from hub. """ pipeline_class_name = MODULAR_PIPELINE_MAPPING.get(self.model_name, ModularPipeline.__name__) diffusers_module = importlib.import_module("diffusers") @@ -1582,7 +1582,7 @@ def __init__( if name in self._component_specs and isinstance(value, (tuple, list)) and len(value) == 2: library, class_name = value component_spec_dict = { - "repo": pretrained_model_name_or_path, + "pretrained_model_name_or_path": pretrained_model_name_or_path, "subfolder": name, "type_hint": (library, class_name), } @@ -1639,7 +1639,7 @@ def from_pretrained( pretrained_model_name_or_path (`str` or `os.PathLike`, optional): Path to a pretrained pipeline configuration. It will first try to load config from `modular_model_index.json`, then fallback to `model_index.json` for compatibility with standard - non-modular repositories. If the repo does not contain any pipeline config, it will be set to None + non-modular repositories. If the pretrained_model_name_or_path does not contain any pipeline config, it will be set to None during initialization. trust_remote_code (`bool`, optional): Whether to trust remote code when loading the pipeline, need to be set to True if you want to create @@ -1809,7 +1809,7 @@ def register_components(self, **kwargs): library, class_name = None, None # extract the loading spec from the updated component spec that'll be used as part of modular_model_index.json config - # e.g. {"repo": "stabilityai/stable-diffusion-2-1", + # e.g. {"pretrained_model_name_or_path": "stabilityai/stable-diffusion-2-1", # "type_hint": ("diffusers", "UNet2DConditionModel"), # "subfolder": "unet", # "variant": None, @@ -2113,7 +2113,7 @@ def load_components(self, names: Optional[Union[List[str], str]] = None, **kwarg **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} - - if potentially override ComponentSpec if passed a different loading field in kwargs, e.g. `repo`, + - if potentially override ComponentSpec if passed a different loading field in kwargs, e.g. `pretrained_model_name_or_path`, `variant`, `revision`, etc. """ @@ -2377,10 +2377,10 @@ def _component_spec_to_dict(component_spec: ComponentSpec) -> Any: - "type_hint": Tuple[str, str] Library name and class name of the component. (e.g. ("diffusers", "UNet2DConditionModel")) - All loading fields defined by `component_spec.loading_fields()`, typically: - - "repo": Optional[str] - The model repository (e.g., "stabilityai/stable-diffusion-xl"). + - "pretrained_model_name_or_path": Optional[str] + The model pretrained_model_name_or_pathsitory (e.g., "stabilityai/stable-diffusion-xl"). - "subfolder": Optional[str] - A subfolder within the repo where this component lives. + A subfolder within the pretrained_model_name_or_path where this component lives. - "variant": Optional[str] An optional variant identifier for the model. - "revision": Optional[str] @@ -2397,11 +2397,11 @@ def _component_spec_to_dict(component_spec: ComponentSpec) -> Any: Example: >>> from diffusers.pipelines.modular_pipeline_utils import ComponentSpec >>> from diffusers import UNet2DConditionModel >>> spec = ComponentSpec( - ... name="unet", ... type_hint=UNet2DConditionModel, ... config=None, ... repo="path/to/repo", ... + ... name="unet", ... type_hint=UNet2DConditionModel, ... config=None, ... pretrained_model_name_or_path="path/to/repo", ... subfolder="subfolder", ... variant=None, ... revision=None, ... default_creation_method="from_pretrained", ... ) >>> ModularPipeline._component_spec_to_dict(spec) { - "type_hint": ("diffusers", "UNet2DConditionModel"), "repo": "path/to/repo", "subfolder": "subfolder", + "type_hint": ("diffusers", "UNet2DConditionModel"), "pretrained_model_name_or_path": "path/to/repo", "subfolder": "subfolder", "variant": None, "revision": None, } """ @@ -2431,10 +2431,10 @@ def _dict_to_component_spec( - "type_hint": Tuple[str, str] Library name and class name of the component. (e.g. ("diffusers", "UNet2DConditionModel")) - All loading fields defined by `component_spec.loading_fields()`, typically: - - "repo": Optional[str] + - "pretrained_model_name_or_path": Optional[str] The model repository (e.g., "stabilityai/stable-diffusion-xl"). - "subfolder": Optional[str] - A subfolder within the repo where this component lives. + A subfolder within the pretrained_model_name_or_path where this component lives. - "variant": Optional[str] An optional variant identifier for the model. - "revision": Optional[str] @@ -2451,10 +2451,10 @@ def _dict_to_component_spec( ComponentSpec: A reconstructed ComponentSpec object. Example: - >>> spec_dict = { ... "type_hint": ("diffusers", "UNet2DConditionModel"), ... "repo": + >>> spec_dict = { ... "type_hint": ("diffusers", "UNet2DConditionModel"), ... "pretrained_model_name_or_path": "stabilityai/stable-diffusion-xl", ... "subfolder": "unet", ... "variant": None, ... "revision": None, ... } >>> ModularPipeline._dict_to_component_spec("unet", spec_dict) ComponentSpec( - name="unet", type_hint=UNet2DConditionModel, config=None, repo="stabilityai/stable-diffusion-xl", + name="unet", type_hint=UNet2DConditionModel, config=None, pretrained_model_name_or_path="stabilityai/stable-diffusion-xl", subfolder="unet", variant=None, revision=None, default_creation_method="from_pretrained" ) """ diff --git a/src/diffusers/modular_pipelines/modular_pipeline_utils.py b/src/diffusers/modular_pipelines/modular_pipeline_utils.py index b15126868634..aa22d963c683 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline_utils.py +++ b/src/diffusers/modular_pipelines/modular_pipeline_utils.py @@ -21,6 +21,7 @@ import torch from ..configuration_utils import ConfigMixin, FrozenDict +from ..loaders.single_file_utils import _validate_single_file_path from ..utils import is_torch_available, logging @@ -80,10 +81,10 @@ class ComponentSpec: type_hint: Type of the component (e.g. UNet2DConditionModel) description: Optional description of the component config: Optional config dict for __init__ creation - repo: Optional repo path for from_pretrained creation - subfolder: Optional subfolder in repo - variant: Optional variant in repo - revision: Optional revision in repo + pretrained_model_name_or_path: Optional pretrained_model_name_or_path path for from_pretrained creation + subfolder: Optional subfolder in pretrained_model_name_or_path + variant: Optional variant in pretrained_model_name_or_path + revision: Optional revision in pretrained_model_name_or_path default_creation_method: Preferred creation method - "from_config" or "from_pretrained" """ @@ -92,7 +93,7 @@ class ComponentSpec: description: Optional[str] = None config: Optional[FrozenDict] = None # YiYi Notes: should we change it to pretrained_model_name_or_path for consistency? a bit long for a field name - repo: Optional[Union[str, List[str]]] = field(default=None, metadata={"loading": True}) + pretrained_model_name_or_path: Optional[Union[str, List[str]]] = field(default=None, metadata={"loading": True}) subfolder: Optional[str] = field(default="", metadata={"loading": True}) variant: Optional[str] = field(default=None, metadata={"loading": True}) revision: Optional[str] = field(default=None, metadata={"loading": True}) @@ -182,7 +183,7 @@ def loading_fields(cls) -> List[str]: @property def load_id(self) -> str: """ - Unique identifier for this spec's pretrained load, composed of repo|subfolder|variant|revision (no empty + Unique identifier for this spec's pretrained load, composed of pretrained_model_name_or_path|subfolder|variant|revision (no empty segments). """ if self.default_creation_method == "from_config": @@ -197,12 +198,12 @@ def decode_load_id(cls, load_id: str) -> Dict[str, Optional[str]]: Decode a load_id string back into a dictionary of loading fields and values. Args: - load_id: The load_id string to decode, format: "repo|subfolder|variant|revision" + load_id: The load_id string to decode, format: "pretrained_model_name_or_path|subfolder|variant|revision" where None values are represented as "null" Returns: Dict mapping loading field names to their values. e.g. { - "repo": "path/to/repo", "subfolder": "subfolder", "variant": "variant", "revision": "revision" + "pretrained_model_name_or_path": "path/to/repo", "subfolder": "subfolder", "variant": "variant", "revision": "revision" } If a segment value is "null", it's replaced with None. Returns None if load_id is "null" (indicating component not created with `load` method). """ @@ -260,33 +261,43 @@ def create(self, config: Optional[Union[FrozenDict, Dict[str, Any]]] = None, **k def load(self, **kwargs) -> Any: """Load component using from_pretrained.""" - # select loading fields from kwargs passed from user: e.g. repo, subfolder, variant, revision, note the list could change + # 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 load_kwargs = {key: passed_loading_kwargs.get(key, getattr(self, key)) for key in self.loading_fields()} - # repo is a required argument for from_pretrained, a.k.a. pretrained_model_name_or_path - repo = load_kwargs.pop("repo", None) - if repo is None: + # pretrained_model_name_or_path is a required argument for from_pretrained, a.k.a. pretrained_model_name_or_path + pretrained_model_name_or_path = load_kwargs.pop("pretrained_model_name_or_path", None) + if pretrained_model_name_or_path is None: raise ValueError( - "`repo` info is required when using `load` method (you can directly set it in `repo` field of the ComponentSpec or pass it as an argument)" + "`pretrained_model_name_or_path` info is required when using `load` method (you can directly set it in `pretrained_model_name_or_path` field of the ComponentSpec or pass it as an argument)" ) + is_single_file = _validate_single_file_path(pretrained_model_name_or_path) + if is_single_file and self.type_hint is None: + raise ValueError("type_hint is required when loading a single file model") if self.type_hint is None: try: from diffusers import AutoModel - component = AutoModel.from_pretrained(repo, **load_kwargs, **kwargs) + component = AutoModel.from_pretrained(pretrained_model_name_or_path, **load_kwargs, **kwargs) except Exception as e: raise ValueError(f"Unable to load {self.name} without `type_hint`: {e}") # update type_hint if AutoModel load successfully self.type_hint = component.__class__ else: + # determine load method + load_method = ( + getattr(self.type_hint, "from_single_file") + if is_single_file + else getattr(self.type_hint, "from_pretrained") + ) + try: - component = self.type_hint.from_pretrained(repo, **load_kwargs, **kwargs) + component = load_method(pretrained_model_name_or_path, **load_kwargs, **kwargs) except Exception as e: raise ValueError(f"Unable to load {self.name} using load method: {e}") - self.repo = repo + self.pretrained_model_name_or_path = pretrained_model_name_or_path for k, v in load_kwargs.items(): setattr(self, k, v) component._diffusers_load_id = self.load_id