make install_path and register_path work; refactor model probing

This commit is contained in:
Lincoln Stein
2023-11-23 23:15:32 -05:00
parent 8c7a7bc897
commit 80bc9be3ab
14 changed files with 1027 additions and 121 deletions

View File

@ -23,7 +23,7 @@ from enum import Enum
from typing import Literal, Optional, Type, Union
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
from typing_extensions import Annotated
from typing_extensions import Annotated, Any, Dict
class InvalidModelConfigException(Exception):
@ -125,7 +125,7 @@ class ModelConfigBase(BaseModel):
validate_assignment=True,
)
def update(self, attributes: dict):
def update(self, attributes: Dict[str, Any]) -> None:
"""Update the object with fields in dict."""
for key, value in attributes.items():
setattr(self, key, value) # may raise a validation error
@ -198,8 +198,6 @@ class MainCheckpointConfig(_CheckpointConfig, _MainConfig):
"""Model config for main checkpoint models."""
type: Literal[ModelType.Main] = ModelType.Main
# Note that we do not need prediction_type or upcast_attention here
# because they are provided in the checkpoint's own config file.
class MainDiffusersConfig(_DiffusersConfig, _MainConfig):