Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Allow embeddings and loras to have different names #6053

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

Allow embeddings and loras to have different names #6053

Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
bbecb99
Allow Embedding and Lora Folder Models to have Different Names
brandonrising Mar 21, 2024
18098cc
For loras/embeddings treat folders the same as single files
brandonrising Mar 21, 2024
befcf77
Fix paths in probe
brandonrising Mar 21, 2024
a7b8bbc
Use model_path defined in probe run
brandonrising Mar 21, 2024
90f32d1
Use correct format type on model
brandonrising Mar 26, 2024
653a8f6
Run ruff
brandonrising Mar 26, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
48 changes: 26 additions & 22 deletions invokeai/backend/model_manager/probe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import json
import os
import re
from pathlib import Path
from typing import Any, Dict, Literal, Optional, Union
Expand Down Expand Up @@ -145,7 +146,8 @@ def probe(
raise InvalidModelConfigException(f"Unhandled combination of {format_type} and {model_type}")

probe = probe_class(model_path)

model_path = probe.model_path
format_type = probe.get_format()
fields["source_type"] = fields.get("source_type") or ModelSourceType.Path
fields["source"] = fields.get("source") or model_path.as_posix()
fields["key"] = fields.get("key", uuid_string())
Expand All @@ -159,7 +161,8 @@ def probe(
fields["description"] = (
fields.get("description") or f"{fields['base'].value} {model_type.value} model {fields['name']}"
)
fields["format"] = fields.get("format") or probe.get_format()
fields["format"] = fields.get("format") or format_type

fields["hash"] = fields.get("hash") or ModelHash(algorithm=hash_algo).hash(model_path)

fields["default_settings"] = fields.get("default_settings")
Expand Down Expand Up @@ -643,17 +646,19 @@ def _guess_name(self) -> str:
return name


class TextualInversionFolderProbe(FolderProbeBase):
def get_format(self) -> ModelFormat:
return ModelFormat.EmbeddingFolder

def get_base_type(self) -> BaseModelType:
path = self.model_path / "learned_embeds.bin"
if not path.exists():
class TextualInversionFolderProbe(TextualInversionCheckpointProbe):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't like how this results in a "folder" model having a path that points to a file...

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, I'm not really sure what value of storing it as a folder even are for these though. I haven't found any lora or embeddings that have other contextual files that affect generations. Just seems to put us in a position where we have to enforce folder structure for no reason to me.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The folder was introduced by hugging face. TIs usually contain two files--one with the weights and another with the name of the concept.

For Lora's the folder is used to give the model a name

def __init__(self, model_path: Path):
files = os.scandir(model_path)
files = [
Path(f.path)
for f in files
if f.is_file() and f.name.endswith((".ckpt", ".pt", ".pth", ".bin", ".safetensors"))
]
if len(files) != 1:
raise InvalidModelConfigException(
f"{self.model_path.as_posix()} does not contain expected 'learned_embeds.bin' file"
f"Unable to determine base type for {model_path}: expected exactly one valid model file, found {[f.name for f in files]}."
)
return TextualInversionCheckpointProbe(path).get_base_type()
super().__init__(files.pop())
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What if the first model file we find isn't the correct one?

I've seen a few TI models hosted on diffusers that have multiple checkpoints in them, presumably representing different stages of trainign.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What if the first model file we find isn't the correct one?

Nevermind that, I misread the length check. The other issue still stands.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't believe any of the checkpoints are actually used for inference but I could be wrong. This only scans the root directory of the folder so as long as there's only 1 model there it loads it fine. It does move it out of the folder though.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What I mean is that this logic would appear to fail probing TIs that have multiple model files in the directory.

Some look like this (struggling to find an example, filenames probably aren't accurate):

  • embedding.pt
  • checkpoint-001.pt
  • checkpoint-002.pt

This is a valid TI and the model loader knows how to load it, but I believe the logic in this PR would see the multiple model files and consider it an invalid model.



class ONNXFolderProbe(PipelineFolderProbe):
Expand Down Expand Up @@ -699,17 +704,16 @@ def get_base_type(self) -> BaseModelType:
return base_model


class LoRAFolderProbe(FolderProbeBase):
def get_base_type(self) -> BaseModelType:
model_file = None
for suffix in ["safetensors", "bin"]:
base_file = self.model_path / f"pytorch_lora_weights.{suffix}"
if base_file.exists():
model_file = base_file
break
if not model_file:
raise InvalidModelConfigException("Unknown LoRA format encountered")
return LoRACheckpointProbe(model_file).get_base_type()
class LoRAFolderProbe(LoRACheckpointProbe):
def __init__(self, model_path: Path):
files = os.scandir(model_path)
files = [Path(f.path) for f in files if f.is_file() and f.name.endswith((".bin", ".safetensors"))]
if len(files) != 1:
raise InvalidModelConfigException(
f"Unable to determine base type for lora {model_path}: expected exactly one valid model file, found {[f.name for f in files]}."
)
model_file = files.pop()
super().__init__(model_file)


class IPAdapterFolderProbe(FolderProbeBase):
Expand Down