ai-simple-engine-runtime-diffusers 0.0.1__tar.gz → 0.0.3__tar.gz

{ai_simple_engine_runtime_diffusers-0.0.1 → ai_simple_engine_runtime_diffusers-0.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ai-simple-engine-runtime-diffusers
-Version: 0.0.1
+Version: 0.0.3
 Summary: AI Simple Engine Runtime Diffusers Module
 License-File: LICENSE
 Author: danialcala94
@@ -9,11 +9,11 @@ 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.2.6,<1.0.0)
+Requires-Dist: ai-simple-engine (>=0.2.8,<1.0.0)
 Requires-Dist: ai-simple-engine-common (>=0.0.3,<1.0.0)
-Requires-Dist: ai-simple-engine-diffusion (>=0.0.2,<1.0.0)
-Requires-Dist: diffusers (>=0.38.0,<1.0.0)
-Requires-Dist: transformers (>=0.0.1,<1.0.0)
+Requires-Dist: ai-simple-engine-diffusion (>=0.0.3,<1.0.0)
+Requires-Dist: diffusers (>=0.38.0,<9999.0.0)
+Requires-Dist: transformers (>=5.13.0,<9999.0.0)
 Description-Content-Type: text/markdown
 
 # AI Simple Engine Runtime Diffusers Module

{ai_simple_engine_runtime_diffusers-0.0.1 → ai_simple_engine_runtime_diffusers-0.0.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ai-simple-engine-runtime-diffusers"
-version = "0.0.1"
+version = "0.0.3"
 description = "AI Simple Engine Runtime Diffusers Module"
 authors = [
     {name = "danialcala94", email = "danielalcalavalera@gmail.com"}
@@ -10,11 +10,11 @@ requires-python = ">=3.10,<3.12"
 
 [tool.poetry.dependencies]
 # Mandatory
-ai-simple-engine = { version = ">=0.2.6,<1.0.0", optional = false }
+ai-simple-engine = { version = ">=0.2.8,<1.0.0", optional = false }
 ai-simple-engine-common = { version = ">=0.0.3,<1.0.0", optional = false }
-ai-simple-engine-diffusion = { version = ">=0.0.2,<1.0.0", optional = false }
-transformers = { version = ">=0.0.1,<1.0.0", optional = false }
-diffusers = { version = ">=0.38.0,<1.0.0", optional = false }
+ai-simple-engine-diffusion = { version = ">=0.0.3,<1.0.0", optional = false }
+transformers = { version = ">=5.13.0,<9999.0.0", optional = false }
+diffusers = { version = ">=0.38.0,<9999.0.0", optional = false }
 # Optional
 
 [tool.poetry]

{ai_simple_engine_runtime_diffusers-0.0.1 → ai_simple_engine_runtime_diffusers-0.0.3}/src/ai_simple_engine_runtime_diffusers/models/executor/diffusers_latent_diffusion_model_executor.py

@@ -4,6 +4,7 @@ from ai_simple_engine.execution.execution_context import ExecutionContext
 from ai_simple_engine_runtime_diffusers.models.executor.prompt_embeddings.prompt_embeddings import PromptEmbeddings
 from ai_simple_engine_runtime_diffusers.models.scheduler.registry.diffusers_scheduler_registry import DiffusersSchedulerRegistry
 from ai_simple_engine.models.loaded_model import LoadedModel
+from ai_simple_engine.types.image import Image
 
 import torch
 
@@ -26,9 +27,9 @@ class DiffusersLatentDiffusionModelExecutor (
         context: ExecutionContext,
         model: LoadedModel,
     ):
-        scheduler_registry = context.services.resolve(DiffusersSchedulerRegistry)
+        scheduler_registry = context.services.get(DiffusersSchedulerRegistry)
 
-        scheduler_class = scheduler_registry.resolve(model.info.scheduler.name)
+        scheduler_class = scheduler_registry.resolve(model.info.scheduler.identifier)
 
         scheduler = scheduler_class.from_config(
             model.instance.scheduler_config,
@@ -109,15 +110,23 @@ class DiffusersLatentDiffusionModelExecutor (
     async def decode_latents(
         self,
         *,
-        model,
+        model: LoadedModel,
         latents
-    ):
-        runtime_model = model.model
+    ) -> Image:
+        runtime_model = model.instance
 
         vae = runtime_model.vae
 
         latents = latents / vae.config.scaling_factor
 
         image = vae.decode(latents).sample
+        """
+        TODO: I don't know if this must be done always
+        or depending or what, but I should pay attention
+        and refactor it.
+        """
+        image = (image / 2 + 0.5).clamp(0, 1)
+        image = image.cpu().permute(0, 2, 3, 1).float().detach().numpy()
+        image = Image(image)
 
         return image

{ai_simple_engine_runtime_diffusers-0.0.1 → ai_simple_engine_runtime_diffusers-0.0.3}/src/ai_simple_engine_runtime_diffusers/models/loader/latent_diffusion_model_loader.py

@@ -9,12 +9,17 @@ from ai_simple_engine.models.loaded_model import LoadedModel
 from transformers import CLIPTokenizer, CLIPTextModel
 from diffusers import UNet2DConditionModel, AutoencoderKL
 
-import torch
-
 
 class LatentDiffusionModelLoader(
     ModelLoader
 ):
+    
+    @property
+    def family(
+        self
+    ) -> str:
+        # TODO: Transform into const
+        return 'latent_diffusion'
 
     async def load(
         self,
@@ -27,7 +32,7 @@ class LatentDiffusionModelLoader(
             path,
             subfolder = 'tokenizer',
             torch_dtype = get_torch_dtype_for(device)
-        ).to(str(device))
+        )#.to(str(device))
 
         text_encoder = CLIPTextModel.from_pretrained(
             path,
@@ -47,11 +52,25 @@ class LatentDiffusionModelLoader(
             torch_dtype = get_torch_dtype_for(device)
         ).to(str(device))
 
+        # TODO: This is if I want its own Scheduler
+        # DDIMScheduler.from_pretrained(
+        #     path,
+        #     subfolder = 'scheduler',
+        #     torch_dtype = get_torch_dtype_for(device)
+        # )#.to(str(device))
+
+        # TODO: We had this before but was failing
         scheduler_config = AutoencoderKL.load_config(
-            path,
+            path + '/scheduler/scheduler_config.json',
             subfolder = 'scheduler',
             torch_dtype = get_torch_dtype_for(device)
-        ).to(str(device))
+        )#.to(str(device))
+
+        # Free memory
+        # TODO: Explain why (?)
+        unet.eval()
+        text_encoder.eval()
+        vae.eval()
 
         runtime_model = DiffusersLatentDiffusionModel(
             tokenizer = tokenizer,

{ai_simple_engine_runtime_diffusers-0.0.1 → ai_simple_engine_runtime_diffusers-0.0.3}/src/ai_simple_engine_runtime_diffusers/models/scheduler/registry/diffusers_scheduler_registry.py

@@ -5,8 +5,10 @@ from diffusers.schedulers.scheduling_utils import SchedulerMixin
 class DiffusersSchedulerRegistry(
     ResolverRegistry[
         str,
-        type[SchedulerMixin]
-    ]
+        type[SchedulerMixin],
+        str
+    ],
+    # Generic[TModel, TExecutor]
 ):
     """
     Class that will register the different

ai_simple_engine_runtime_diffusers-0.0.3/src/ai_simple_engine_runtime_diffusers/models/scheduler/utils.py

@@ -0,0 +1,22 @@
+"""
+TODO: Maybe rename and move this module.
+"""
+from diffusers.schedulers import (
+    DDIMScheduler,
+    EulerDiscreteScheduler,
+    EulerAncestralDiscreteScheduler,
+    DPMSolverMultistepScheduler,
+    FlowMatchEulerDiscreteScheduler
+)
+
+SCHEDULERS = {
+    'ddim': DDIMScheduler,
+    'euler': EulerDiscreteScheduler,
+    'euler_a': EulerAncestralDiscreteScheduler,
+    'dpmpp_2m': DPMSolverMultistepScheduler,
+    'flow_match_euler': FlowMatchEulerDiscreteScheduler
+}
+"""
+The list including all the schedulers we are
+accepting.
+"""

{ai_simple_engine_runtime_diffusers-0.0.1 → ai_simple_engine_runtime_diffusers-0.0.3}/src/ai_simple_engine_runtime_diffusers/plugin.py

@@ -1,11 +1,15 @@
 from ai_simple_engine_runtime_diffusers.models.loader.latent_diffusion_model_loader import LatentDiffusionModelLoader
-from ai_simple_engine_runtime_diffusers.models.scheduler.utils import SCHEDULERS
+from ai_simple_engine_runtime_diffusers.models.executor.noise.pytorch_noise_generator import TorchNoiseGenerator
 from ai_simple_engine_runtime_diffusers.models.scheduler.registry.diffusers_scheduler_registry import DiffusersSchedulerRegistry
+from ai_simple_engine.models.executor.registry.family_model_executor_registry import FamilyModelExecutorRegistry
+from ai_simple_engine_runtime_diffusers.models.executor.diffusers_latent_diffusion_model_executor import DiffusersLatentDiffusionModelExecutor
+from ai_simple_engine_runtime_diffusers.models.scheduler.utils import SCHEDULERS
+from ai_simple_engine_diffusion.noise.noise_generator_abstract import NoiseGenerator
 from ai_simple_engine.engine_builder import EngineBuilder
 from ai_simple_engine.plugins.plugin import Plugin
 
 
-class DiffusersRuntimeMusicgenPlugin(
+class DiffusersRuntimePlugin(
     Plugin
 ):
     """
@@ -19,10 +23,15 @@ class DiffusersRuntimeMusicgenPlugin(
         self,
         builder: EngineBuilder
     ):
-        
+        (
+            builder.add_model_loader(LatentDiffusionModelLoader())
+        )
 
         (
-            builder.add_model_loader(LatentDiffusionModelLoader)
+            builder.add_service(
+                NoiseGenerator,
+                TorchNoiseGenerator()
+            )
         )
 
         # TODO: Register the SchedulerRegistry
@@ -36,34 +45,17 @@ class DiffusersRuntimeMusicgenPlugin(
             scheduler_spec.identifier
         )
         """
+        model_executor_registry = builder.get_or_add_service(FamilyModelExecutorRegistry)
+
+        model_executor_registry.register(
+            'latent_diffusion',
+            DiffusersLatentDiffusionModelExecutor()
+        )
+
         scheduler_registry = builder.get_or_add_service(DiffusersSchedulerRegistry)
 
         for identifier, cls in SCHEDULERS.items():
             scheduler_registry.register(
                 identifier,
                 cls
-            )
-        
-        # scheduler_registry.register(
-        #     'euler',
-        #     EulerDiscreteScheduler
-        # )
-
-        # TODO: This below was the 'musicgen'
-        # (
-        #     builder
-        #     .add_model_loader(MusicgenModelLoader())
-        # )
-
-        # """
-        # We obtain the registry that handles the
-        # model executors by the model's family,
-        # and we register our specific model
-        # executor that uses transformers.
-        # """
-        # registry = builder.get_or_add_service(FamilyModelExecutorRegistry)
-
-        # registry.register(
-        #     MUSICGEN_MODEL_FAMILY,
-        #     TransformersMusicgenModelExecutor()
-        # )
+            )

ai_simple_engine_runtime_diffusers-0.0.1/src/ai_simple_engine_runtime_diffusers/models/scheduler/utils.py

@@ -1,44 +0,0 @@
-"""
-TODO: Maybe rename and move this module.
-"""
-from diffusers.schedulers import (
-    DDIMScheduler,
-    EulerDiscreteScheduler,
-    EulerAncestralDiscreteScheduler,
-    DPMSolverMultistepScheduler,
-    FlowMatchEulerDiscreteScheduler
-)
-
-SCHEDULERS = {
-    'ddim': DDIMScheduler,
-    'euler': EulerDiscreteScheduler,
-    'euler_a': EulerAncestralDiscreteScheduler,
-    'dpmpp_2m': DPMSolverMultistepScheduler,
-    'flow_match_euler': FlowMatchEulerDiscreteScheduler
-}
-"""
-The list including all the schedulers we are
-accepting.
-"""
-
-
-def get_scheduler_class(
-    scheduler: str,
-    model: LoadedModel
-) -> Union[SchedulerMixin, None]:
-    scheduler_class = SCHEDULERS.get(scheduler, None)
-
-    if scheduler_class is None:
-        raise Exception(f'The scheduler "{scheduler}" is not accepted by our system.')
-    
-    return scheduler_class
-
-    scheduler_cls = registry.resolve(model.info.scheduler.identifier)
-
-    scheduler = scheduler_cls.from_config(...)
-
-    return scheduler
-
-    return scheduler_cls.from_config(
-        model.model.scheduler.config
-    )