zodiac.streams.class_stream

  1#  # # <!-- // /*  SPDX-License-Identifier: MPL-2.0*/ -->
  2#  # # <!-- // /*  d a r k s h a p e s */ -->
  3
  4from typing import List, Tuple, Callable, Union, Any, Generator
  5from zodiac.providers.registry_entry import RegistryEntry
  6from zodiac.providers.constants import MIR_DB, VERSIONS_CONFIG, ChipType
  7
  8
  9async def ancestor_data(mir_tag_or_registry_entry: RegistryEntry | list, field_name: str = "pkg") -> Generator:
 10    """Trace lineage of a model for the specified field \n
 11    :param registry_entry: RegistryEntry for the model that needs to be traced
 12    :param field_name: The name of the database field containing the data sought
 13    :return: A generator populated with matching data fields"""
 14    mir_db = MIR_DB.database
 15    if isinstance(mir_tag_or_registry_entry, RegistryEntry):
 16        mir_prefix = mir_tag_or_registry_entry.mir[0]
 17    else:
 18        mir_prefix = mir_tag_or_registry_entry[0]
 19    base_fields = [
 20        "diffusers",
 21        "*",
 22    ]  # "prior"
 23    return [mir_db[mir_prefix][x].get(field_name) for x in base_fields if mir_db[mir_prefix].get(x, {}).get(field_name, {})]
 24
 25
 26async def best_package(pkg_data: RegistryEntry | dict[int | str, Any], ready_list: list[tuple[ChipType]] = ChipType._show_ready()) -> tuple[str]:
 27    """Identify the best package based on model data and package sets.\n
 28    :param mir_db_pkg: Dictionary containing package data to match
 29    :param ready_pkg_types: List of priority package processors to evaluate
 30    :return: Tuple containing (class name, package type) if match found, otherwise None"""
 31
 32    if isinstance(pkg_data, RegistryEntry):
 33        pkg_loop: list = await ancestor_data(pkg_data)
 34        # print(pkg_loop)
 35        pkg_loop.insert(0, pkg_data.modules | pkg_loop[0])
 36    else:
 37        pkg_loop = [pkg_data]  # await ancestor_data(pkg_data)  # normalize to list
 38    for processor in ready_list:
 39        for pkg_type in processor[2]:
 40            if pkg_type.value[0]:  # Determine if the package is available
 41                package_name = pkg_type.value[1].lower()
 42                for index, data in next(iter(pkg_loop)).items():
 43                    if package_name in data:
 44                        return (index, data, pkg_type)
 45
 46
 47async def find_package(entry: RegistryEntry = None, mir_entry: list[str] | None = None) -> Tuple[str]:
 48    """Look up class and package in MIR from RegistryEntry.\n
 49    :param entry: A RegistryEntry object containing MIR (Model Identifier Resource) details.
 50    :return: A tuple containing the class name of the package and its type if found; otherwise, None.
 51    :raises: AttributeError: If PkgType or ChipType classes are not properly defined."""
 52    import re
 53
 54    mir_base = entry.mir[0] if not mir_entry else mir_entry[0]
 55    mir_comp = entry.mir[1] if not mir_entry else mir_entry[1]
 56    mir_ids = [mir_comp, "diffusers", "*"]
 57    suffixes = VERSIONS_CONFIG.get("suffixes")
 58    if suffixes:
 59        for compatibility, model_data in MIR_DB.database[mir_base].items():
 60            package_key = model_data.get("pkg")
 61            if package_key and (any(re.match(pattern, compatibility) for pattern in suffixes) or compatibility in mir_ids):  # Fallback to CPU packages if no match found in GPU packages
 62                package_data = await best_package(package_key)
 63                if package_data:
 64                    return package_data
 65    return
 66
 67
 68async def stage_class(class_object: Callable) -> List[Tuple[Union[str, Callable]]]:
 69    """Returns a tuple of data for each sub-class of a module
 70    :param class_obj: The class item to inspect.
 71    ex:('diffusers', 'models.autoencoders.autoencoder_kl', 'AutoencoderKL', <class 'diffusers.models.autoencoders.autoencoder_kl.AutoencoderKL'>),"""
 72    import typing
 73
 74    pipe_args = typing.get_type_hints(class_object.__init__).values()
 75    sub_classes = [
 76        (*pipe.__module__.split(".", 1), pipe.__name__, pipe)  #
 77        for pipe in pipe_args  #
 78        if "builtins" not in pipe.__module__ and "typing" not in pipe.__module__
 79    ]
 80    return sub_classes
 81
 82
 83# async def show_transformers_tasks(class_name: str) -> List[str]:
 84#     """Retrieves a list of task classes associated with a specified transformer class.\n
 85#     :param class_name: The name of the transformer class to inspect.
 86#     :param pkg_type: The dependency for the module
 87#     :return: A list of task classes associated with the specified transformer.
 88#     """
 89#     class_obj: Callable = make_callable(class_name, PkgType.TRANSFORMERS.value[1].lower())
 90#     class_module: Callable = make_callable(*class_obj.__module__.split(".", 1)[-1:], class_obj.__module__.split(".", 1)[0])
 91#     task_classes = getattr(class_module, "__all__")
 92#     return task_classes
 93
 94
 95# from mir.mappers import make_callable
 96# from zodiac.providers.constants import MIR_DB
 97# from zodiac.class_stream import lookup_package
 98# from zodiac.class_stream import trace_class
 99
100# model_data = lookup_package(entry)
101# package_data = [content for content in model_data.get("pkg").values() if next(iter(content)) in ["diffusers", "transformers"]]
102# class_data = trace_class(make_callable(package_data[0], "diffusers"))
103# # [terminal_gen(data[3],getattr(PkgType,data[0].upper())) for data in class_data]