ai-simple-engine-transformers 0.0.1__py3-none-any.whl

ai_simple_engine_transformers/__init__.py

@@ -0,0 +1,10 @@
+"""
+The Transformers module of our AI Simple Engine.
+"""
+
+
+"""
+TODO: Check these things:
+- Should the 'InstalledModel' have the 'ModelSpec'
+inside so we can access to the 'family' and that (?)
+"""

ai_simple_engine_transformers/model_loaders/musicgen/base.py

@@ -0,0 +1,42 @@
+from ai_simple_engine_transformers.model_loaders.musicgen.loaded_musicgen_model import LoadedMusicGenModel
+from ai_simple_engine.models.loaded_model import LoadedModel
+from ai_simple_engine.models.loaders.abstract import ModelLoader
+from ai_simple_engine.models.installed_model import InstalledModel
+from ai_simple_engine.device.base import Device
+from transformers import AutoProcessor, MusicgenForConditionalGeneration
+
+
+class MusicGenLoader(
+    ModelLoader
+):
+
+    def is_supported(
+        self,
+        family: str
+    ) -> bool:
+        return family == 'musicgen'
+
+    async def load(
+        self,
+        model: InstalledModel,
+        *,
+        device: Device
+    ) -> LoadedModel[LoadedMusicGenModel]:
+        processor = AutoProcessor.from_pretrained(model.path)
+        network = MusicgenForConditionalGeneration.from_pretrained(model.path)
+
+        network.to(device.identifier)
+
+        return LoadedModel(
+            installed_model = model,
+            instance = LoadedMusicGenModel(
+                processor = processor,
+                network = network
+            )
+        )
+
+    async def unload(
+        self,
+        model: LoadedModel
+    ) -> None:
+        del model

ai_simple_engine_transformers/model_loaders/musicgen/loaded_musicgen_model.py

@@ -0,0 +1,9 @@
+from transformers import AutoProcessor, MusicgenForConditionalGeneration
+from dataclasses import dataclass
+
+
+@dataclass(frozen = True)
+class LoadedMusicGenModel:
+
+    processor: AutoProcessor
+    network: MusicgenForConditionalGeneration

ai_simple_engine_transformers/operations/music/__init__.py

@@ -0,0 +1,3 @@
+"""
+TODO: Maybe I don't need the folders
+"""

ai_simple_engine_transformers/operations/music/generate_music.py

@@ -0,0 +1,76 @@
+from ai_simple_engine_transformers.runners.musicgen_runner import MusicGenRunner
+from ai_simple_engine_transformers.types.data_type import RESOURCE, AUDIO
+from ai_simple_engine.graph.operation.base import Operation
+from ai_simple_engine.graph.input import Input
+from ai_simple_engine.graph.output import Output
+from ai_simple_engine.types.data_type import STRING
+from ai_simple_engine.types.audio import Audio
+
+
+class GenerateMusic(
+    Operation
+):
+
+    loaded_model: 'LoadedModel[LoadedMusicGenModel]' = Input(RESOURCE)
+    prompt = Input(STRING)
+    max_new_tokens = Input(
+        int,
+        default = 256
+    )
+    guidance_scale = Input(
+        float,
+        default = 3.0
+    )
+    audio = Output(AUDIO)
+
+    async def execute(
+        self,
+        context
+    ):
+        """
+        TODO: We will implement an OperationFactory to
+        make this a 'self._runner' that is automatically
+        registered.
+        """
+        runner = context.runners.get(MusicGenRunner)
+
+        audio = await runner.generate(
+            model = self.loaded_model.instance,
+            prompt = self.prompt,
+            max_new_tokens = self.max_new_tokens,
+            guidance_scale = self.guidance_scale
+        )
+
+        return {
+            'audio': audio
+        }
+
+
+        loaded_model = self.model.instance
+        processor = loaded_model.processor
+        network = loaded_model.network
+
+        inputs = processor(
+            text = [self.prompt],
+            return_tensors = 'pt'
+        )
+
+        inputs = {
+            key: value.to(network.device)
+            for key, value in inputs.items()
+        }
+
+        output = network.generate(
+            **inputs,
+            max_new_tokens = self.max_new_tokens,
+            guidance_scale = self.guidance_scale
+        )
+
+        audio = output.cpu().numpy()
+
+        return {
+            'audio': Audio(
+                samples = audio[0, 0],
+                sample_rate = 32000
+            )
+        }

ai_simple_engine_transformers/plugin.py

@@ -0,0 +1,18 @@
+from ai_simple_engine.plugins.plugin import Plugin
+from ai_simple_engine_transformers.providers.model_loaders_provider import TransformersModelLoaderProvider
+from ai_simple_engine_transformers.providers.model_repository_provider import TransformersModelRepositoryProvider
+from ai_simple_engine_transformers.providers.operation_provider import TransformersOperationProvider
+
+
+class TransformersPlugin(
+    Plugin
+):
+
+    def providers(
+        self
+    ):
+        return [
+            TransformersModelRepositoryProvider(),
+            TransformersModelLoaderProvider(),
+            TransformersOperationProvider()
+        ]

ai_simple_engine_transformers/providers/__init__.py

@@ -0,0 +1,6 @@
+"""
+The providers that will be used by the
+engine builder to register the different
+elements provided in order to be able to
+work with the plugin.
+"""

ai_simple_engine_transformers/providers/data_types_provider.py

@@ -0,0 +1,18 @@
+from ai_simple_engine.engine_builder import EngineBuilder
+from ai_simple_engine.plugins.providers.abstract import PluginProvider
+from ai_simple_engine_transformers.types.data_type import MUSICGEN_MODEL
+
+
+class TransformersDataTypeProvider(
+    PluginProvider
+):
+
+    def register(
+        self,
+        builder: EngineBuilder
+    ):
+        (
+            builder
+            .add_data_type(MUSICGEN_MODEL)
+        )
+        

ai_simple_engine_transformers/providers/model_loaders_provider.py

@@ -0,0 +1,17 @@
+from ai_simple_engine.engine_builder import EngineBuilder
+from ai_simple_engine.plugins.providers.abstract import PluginProvider
+from ai_simple_engine_transformers.model_loaders.musicgen.base import MusicGenLoader
+
+
+class TransformersModelLoaderProvider(
+    PluginProvider
+):
+
+    def register(
+        self,
+        builder: EngineBuilder
+    ):
+        (
+            builder
+            .add_model_loader(MusicGenLoader())
+        )

ai_simple_engine_transformers/providers/model_repository_provider.py

@@ -0,0 +1,18 @@
+from ai_simple_engine.engine_builder import EngineBuilder
+from ai_simple_engine.plugins.providers.abstract import PluginProvider
+from ai_simple_engine_transformers.repositories.huggingface.huggingface_provider import HuggingFaceProvider
+
+
+class TransformersModelRepositoryProvider(
+    PluginProvider
+):
+
+    def register(
+        self,
+        builder: EngineBuilder
+    ):
+        (
+            builder
+            .add_model_provider(HuggingFaceProvider())
+        )
+        

ai_simple_engine_transformers/providers/operation_provider.py

@@ -0,0 +1,21 @@
+from ai_simple_engine.engine_builder import EngineBuilder
+from ai_simple_engine.plugins.providers.abstract import PluginProvider
+from ai_simple_engine.graph.operation.install_model import InstallModel
+from ai_simple_engine.graph.operation.acquire_model import AcquireModel
+from ai_simple_engine_transformers.operations.music.generate_music import GenerateMusic
+
+
+class TransformersOperationProvider(
+    PluginProvider
+):
+
+    def register(
+        self,
+        builder: EngineBuilder
+    ):
+        (
+            builder
+            .add_operation(InstallModel)
+            .add_operation(AcquireModel)
+            .add_operation(GenerateMusic)
+        )

ai_simple_engine_transformers/repositories/huggingface/huggingface_client.py

@@ -0,0 +1,23 @@
+from huggingface_hub import snapshot_download
+from pathlib import Path
+
+
+# TODO: Maybe I need a 'Client' abstract with the
+# 'download' as an abstract method
+class HuggingFaceClient:
+
+    def download(
+        self,
+        repo_id: str,
+        *,
+        revision: str | None = None,
+        local_dir: Path | None = None
+    ) -> Path:
+        path = snapshot_download(
+            repo_id = repo_id,
+            revision = revision,
+            local_dir = local_dir,
+            local_dir_use_symlinks = False
+        )
+
+        return Path(path)

ai_simple_engine_transformers/repositories/huggingface/huggingface_provider.py

@@ -0,0 +1,102 @@
+from ai_simple_engine_transformers.repositories.huggingface.huggingface_client import HuggingFaceClient
+from ai_simple_engine.models.installed_model import InstalledModel
+from ai_simple_engine.models.providers.abstract import ModelProvider
+from ai_simple_engine.models.model_spec import ModelSpec
+from ai_simple_engine.settings.engine_settings import EngineSettings
+from pathlib import Path
+from typing import Union
+
+
+class HuggingFaceProvider(
+    ModelProvider
+):
+
+    def __init__(
+        self,
+        client: Union[HuggingFaceClient, None] = None,
+    ):
+        self._client = (
+            client or
+            HuggingFaceClient()
+        )
+
+        self._settings: Union[EngineSettings, None] = None
+        
+        # self._models_directory = (
+        #     models_directory
+        #     if models_directory is not None
+        #     # TODO: Change this path
+        #     else Path.home() / '.ai_simple_engine' / 'models'
+        # )
+
+    @property
+    def name(
+        self
+    ) -> str:
+        return 'huggingface'
+
+    async def install(
+        self,
+        spec: ModelSpec
+    ) -> InstalledModel:
+        path = self._client.download(
+            repo_id = spec.identifier,
+            revision = spec.revision,
+            local_dir = self._local_directory(spec)
+        )
+
+        return InstalledModel(
+            id = spec.identifier,
+            # family = spec.family,
+            path = path,
+            revision = spec.revision
+        )
+
+    def is_installed(
+        self,
+        spec: ModelSpec
+    ) -> bool:
+        """
+        Check if the model with the `spec` given is
+        installed or not.
+        """
+        return self._local_directory(spec).exists()
+
+    def get_installed_model(
+        self,
+        spec: ModelSpec
+    ) -> InstalledModel:
+        path = self._local_directory(spec)
+
+        if not path.exists():
+            raise FileNotFoundError(
+                f'Model "{spec.identifier}" is not installed.'
+            )
+
+        return InstalledModel(
+            id = spec.identifier,
+            family = spec.family,
+            path = path,
+            revision = spec.revision
+        )
+
+    def _local_directory(
+        self,
+        # TODO: Remove this spec
+        spec: ModelSpec
+    ) -> Path:
+        # TODO: The '.configure(plugin_context)' should be called
+        return (
+            self._settings.models_directory
+            / spec.family
+            / spec.identifier
+        )
+
+        return self._settings.models_directory
+        
+        path = self._models_directory / spec.family / spec.identifier
+
+        if spec.revision:
+            path /= spec.revision
+
+        return path

ai_simple_engine_transformers/runners/__init__.py

@@ -0,0 +1,12 @@
+"""
+A runner is the class that will include the
+code to perform a real operation. The logic,
+the core.
+
+What you should implement:
+- Abstract class to determine what should the
+`generate` method expect as parameters.
+- Real implementation of the logic that must
+be actually applied to obtain the expected
+results.
+"""

ai_simple_engine_transformers/runners/abstract/musicgen_runner_abstract.py

@@ -0,0 +1,21 @@
+from ai_simple_engine_transformers.model_loaders.musicgen.loaded_musicgen_model import LoadedMusicGenModel
+from ai_simple_engine.runners.abstract import Runner
+from ai_simple_engine.types.audio import Audio
+from abc import ABC, abstractmethod
+
+
+class MusicGenRunnerAbstract(
+    Runner,
+    ABC
+):
+
+    @abstractmethod
+    async def generate(
+        self,
+        model: LoadedMusicGenModel,
+        prompt: str,
+        *,
+        max_new_tokens: int,
+        guidance_scale: float
+    ) -> Audio:
+        ...

ai_simple_engine_transformers/runners/musicgen_runner.py

@@ -0,0 +1,41 @@
+from ai_simple_engine_transformers.runners.abstract.musicgen_runner_abstract import MusicGenRunnerAbstract
+from ai_simple_engine.types.audio import Audio
+
+
+class MusicGenRunner(
+    MusicGenRunnerAbstract
+):
+
+    async def generate(
+        self,
+        model,
+        prompt,
+        *,
+        max_new_tokens,
+        guidance_scale
+    ) -> Audio:
+        processor = model.processor
+        network = model.network
+
+        inputs = processor(
+            text = [prompt],
+            return_tensors = 'pt'
+        )
+
+        inputs = {
+            k: v.to(network.device)
+            for k, v in inputs.items()
+        }
+
+        output = network.generate(
+            **inputs,
+            max_new_tokens = max_new_tokens,
+            guidance_scale = guidance_scale
+        )
+
+        samples = output.cpu().numpy()[0, 0]
+
+        return Audio(
+            samples = samples,
+            sample_rate = 32000
+        )

ai_simple_engine_transformers/types/data_type.py

@@ -0,0 +1,26 @@
+from ai_simple_engine.types.data_type import DataType, RESOURCE
+from ai_simple_engine.resources.resource_handle import ResourceHandle
+
+import torch
+
+
+TENSOR = DataType(
+    name = 'Tensor',
+    runtime_type = torch.Tensor,
+    parent = None
+)
+LATENT = DataType(
+    name = 'Latent',
+    runtime_type = torch.Tensor,
+    parent = TENSOR
+)
+MUSICGEN_MODEL = DataType(
+    name = 'MusicGenModel',
+    runtime_type = ResourceHandle,
+    parent = RESOURCE
+)
+# WHISPER_MODEL = DataType(
+#     name = 'WhisperMode',
+#     runtime_type = ResourceHandle,
+#     parent = RESOURCE
+# )

ai_simple_engine_transformers-0.0.1.dist-info/METADATA

@@ -0,0 +1,20 @@
+Metadata-Version: 2.4
+Name: ai-simple-engine-transformers
+Version: 0.0.1
+Summary: AI Simple Engine Transformers Module
+License-File: LICENSE
+Author: danialcala94
+Author-email: danielalcalavalera@gmail.com
+Requires-Python: >=3.10,<3.12
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Requires-Dist: ai-simple-engine (>=0.0.1,<1.0.0)
+Requires-Dist: huggingface-hub (>=1.21.0,<9999.0.0)
+Requires-Dist: torch (>=2.12.1,<9999.0.0)
+Requires-Dist: transformers (>=5.12.1,<9999.0.0)
+Description-Content-Type: text/markdown
+
+# AI Simple Engine Transformers Module
+
+The Transformers module for the AI Simple Engine `ai-simple-engine`.

ai_simple_engine_transformers-0.0.1.dist-info/RECORD

@@ -0,0 +1,28 @@
+ai_simple_engine_transformers/__init__.py,sha256=RhXJ6s5ed_wiymG6YqMi5OuFrbi-B-sLZ8QnUXQC7LQ,205
+ai_simple_engine_transformers/model_loaders/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ai_simple_engine_transformers/model_loaders/musicgen/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ai_simple_engine_transformers/model_loaders/musicgen/base.py,sha256=I0jzyf-DMYQQJaWwBy3BWqQhYNavudndDQckJ1aEz4c,1239
+ai_simple_engine_transformers/model_loaders/musicgen/loaded_musicgen_model.py,sha256=1f8XWq1p27UWGaT1BCNF8nmfcZpdwLae1ejlqBgVPVU,245
+ai_simple_engine_transformers/operations/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ai_simple_engine_transformers/operations/music/__init__.py,sha256=S6y4Ga1R--d1HuGFK-vvwzous5Bgl26kMohWXVRiuQM,46
+ai_simple_engine_transformers/operations/music/generate_music.py,sha256=ydhL6fR_rYGoCDWEcuJIjczcjyI2nuN-et-YNiL_Lcg,2046
+ai_simple_engine_transformers/plugin.py,sha256=_rToq0bIpSMY0qqEvMSur_iFZX8w6I0kmi51UCsU9ug,640
+ai_simple_engine_transformers/providers/__init__.py,sha256=FA0L0Fbn6P20bW3IEde3t-4-QZt3vjghJL_YIJv2kMQ,155
+ai_simple_engine_transformers/providers/data_types_provider.py,sha256=pru8izdgimvnvqUladZ1wWNG87fwWkbTPVWp12BhxlQ,441
+ai_simple_engine_transformers/providers/model_loaders_provider.py,sha256=wwJfOpPmGIs5JI1VZaqR5niOEYF9pNqj8lIsYMTzKmU,451
+ai_simple_engine_transformers/providers/model_repository_provider.py,sha256=rClroigTbcS-AOhCsM2AQt9A6o8BnUTegeOcTpHznxc,495
+ai_simple_engine_transformers/providers/operation_provider.py,sha256=mqo8DHl-tPmquqscpIBiby1lnrFQAj7xPKF_iKmQYwo,676
+ai_simple_engine_transformers/repositories/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ai_simple_engine_transformers/repositories/huggingface/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ai_simple_engine_transformers/repositories/huggingface/huggingface_client.py,sha256=3g1pnMF-v4D64X9oUBIea8uwgYFEwbg473-AUw0mAaM,573
+ai_simple_engine_transformers/repositories/huggingface/huggingface_provider.py,sha256=AzYlcn-UB2l49ZTcpsOgMrt9hP-_JxYd3auZti1jOMM,2769
+ai_simple_engine_transformers/runners/__init__.py,sha256=0S7Q4W0avUh4YWt4eeTX5M3F6741zKyf0g563N_Im-4,328
+ai_simple_engine_transformers/runners/abstract/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ai_simple_engine_transformers/runners/abstract/musicgen_runner_abstract.py,sha256=EuU1Owexp_fuMAtVkcACexDrZYoQ-lQdZOpekCNwf30,531
+ai_simple_engine_transformers/runners/musicgen_runner.py,sha256=YtjmV9Vnc342tRPOYSkh8c0snOqFvl8Tux8xHZWMXLI,965
+ai_simple_engine_transformers/types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ai_simple_engine_transformers/types/data_type.py,sha256=aJdx0uxOBCTbUQK3sjKuMfkzsWkW4upwEXYYAMyFr98,597
+ai_simple_engine_transformers-0.0.1.dist-info/licenses/LICENSE,sha256=6kbiFSfobTZ7beWiKnHpN902HgBx-Jzgcme0SvKqhKY,1091
+ai_simple_engine_transformers-0.0.1.dist-info/METADATA,sha256=nTzFLo2Gz3SMMg8xwM9LzGHfyNVf5ALVkPTBb00SIac,727
+ai_simple_engine_transformers-0.0.1.dist-info/WHEEL,sha256=M5asmiAlL6HEcOq52Yi5mmk9KmTVjY2RDPtO4p9DMrc,88
+ai_simple_engine_transformers-0.0.1.dist-info/RECORD,,

ai_simple_engine_transformers-0.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: poetry-core 2.2.0
+Root-Is-Purelib: true
+Tag: py3-none-any

ai_simple_engine_transformers-0.0.1.dist-info/licenses/LICENSE

@@ -0,0 +1,19 @@
+Copyright (c) 2018 The Python Packaging Authority
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.