ai-simple-engine-runtime-diffusers 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

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/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/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/models/scheduler/utils.py

@@ -19,26 +19,4 @@ SCHEDULERS = {
 """
 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
-    )
+"""

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.dist-info → ai_simple_engine_runtime_diffusers-0.0.3.dist-info}/METADATA

@@ -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.dist-info → ai_simple_engine_runtime_diffusers-0.0.3.dist-info}/RECORD

@@ -1,7 +1,7 @@
 ai_simple_engine_runtime_diffusers/__init__.py,sha256=oeCwDniRRNaeu1i8oSoIGl24B_le57R-tjJfbQIcaTA,101
 ai_simple_engine_runtime_diffusers/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ai_simple_engine_runtime_diffusers/models/executor/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ai_simple_engine_runtime_diffusers/models/executor/diffusers_latent_diffusion_model_executor.py,sha256=-kVvNpBZLbcWuHiXo6cj3-Y6mVXCn9YIN9S_7opKbj0,3701
+ai_simple_engine_runtime_diffusers/models/executor/diffusers_latent_diffusion_model_executor.py,sha256=mBKcpDmnSo-AroNNaMlQq64aqpynbIm_ULY2aaXi8KE,4093
 ai_simple_engine_runtime_diffusers/models/executor/noise/__init__.py,sha256=gtBak1TwMesgIsuc56HdYmCwrTjdysO68FO6XzAs-ME,69
 ai_simple_engine_runtime_diffusers/models/executor/noise/pytorch_noise_generator.py,sha256=kFL42BiG2odZp7hAA-7FZakihakOzc8BpYorysDvWx0,508
 ai_simple_engine_runtime_diffusers/models/executor/prompt_embeddings/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -10,17 +10,17 @@ ai_simple_engine_runtime_diffusers/models/executor/text_encoder_model_executor_a
 ai_simple_engine_runtime_diffusers/models/executor/unet_model_executor_abstract.py,sha256=xj5Fw8J1jAwJZRoE5dEnH4Ewf92jIyEh6BJEOHOY1n0,1004
 ai_simple_engine_runtime_diffusers/models/executor/vae_model_executor_abstract.py,sha256=t3j4jYHKbo7bSjMhLLOJYsY2IZm7uQGXXK-c7QDCTv8,1045
 ai_simple_engine_runtime_diffusers/models/loader/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ai_simple_engine_runtime_diffusers/models/loader/latent_diffusion_model_loader.py,sha256=fROYlfzuJCqmneCm7xr0OY2PRtKjVeL5JCVi1c-SWP0,3057
+ai_simple_engine_runtime_diffusers/models/loader/latent_diffusion_model_loader.py,sha256=C-vENcRTMk9__Es_TdP38j3avsqK2olIFtXXmC_Tcao,3645
 ai_simple_engine_runtime_diffusers/models/runtime/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ai_simple_engine_runtime_diffusers/models/runtime/diffusers_latent_diffusion_model.py,sha256=UTgR2mWW3rDtJXyRATpR_lfcWS6wMugeKHlvKneXTbw,478
 ai_simple_engine_runtime_diffusers/models/runtime/utils.py,sha256=niSrKW6RdZwtRPQVEFCArTyYdGgWWk-rlFSd2ledSPo,347
 ai_simple_engine_runtime_diffusers/models/scheduler/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ai_simple_engine_runtime_diffusers/models/scheduler/diffusers_scheduler.py,sha256=TiirOtBVX9Op8VeYEjmh3F-rxCiJoEkM6ezyKFjQ_74,896
 ai_simple_engine_runtime_diffusers/models/scheduler/registry/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ai_simple_engine_runtime_diffusers/models/scheduler/registry/diffusers_scheduler_registry.py,sha256=i9nAd_T30oFFT65awnIw_Z6cDtctmN1NWxkQArXDugY,516
-ai_simple_engine_runtime_diffusers/models/scheduler/utils.py,sha256=VSVeypJkrwaCTRZyYrmscE8g9_pufQ-AlZAhUD_3OKE,1111
-ai_simple_engine_runtime_diffusers/plugin.py,sha256=rPl2BEW9ow6R0KouaFvvBpetG5Opz92E-Rq-6xH5iLc,2101
-ai_simple_engine_runtime_diffusers-0.0.1.dist-info/licenses/LICENSE,sha256=6kbiFSfobTZ7beWiKnHpN902HgBx-Jzgcme0SvKqhKY,1091
-ai_simple_engine_runtime_diffusers-0.0.1.dist-info/METADATA,sha256=p_wUk1JdE8bWcJGNq3m2tDWCbvaN_N2IptU7HbU3vX8,834
-ai_simple_engine_runtime_diffusers-0.0.1.dist-info/WHEEL,sha256=M5asmiAlL6HEcOq52Yi5mmk9KmTVjY2RDPtO4p9DMrc,88
-ai_simple_engine_runtime_diffusers-0.0.1.dist-info/RECORD,,
+ai_simple_engine_runtime_diffusers/models/scheduler/registry/diffusers_scheduler_registry.py,sha256=FOuIAAwNmi3_jB25ICOlCPbD7U1WlOcGFZT3rhql5cY,565
+ai_simple_engine_runtime_diffusers/models/scheduler/utils.py,sha256=tIzwp7fy0Y1pE7XcLJU0dEp4XH2TMxBzY7SnvyfwITM,554
+ai_simple_engine_runtime_diffusers/plugin.py,sha256=o_5OJSwHj_4Qy1LW-Y_JRIP3-LotSOEUVgB3LdnK3mU,2205
+ai_simple_engine_runtime_diffusers-0.0.3.dist-info/licenses/LICENSE,sha256=6kbiFSfobTZ7beWiKnHpN902HgBx-Jzgcme0SvKqhKY,1091
+ai_simple_engine_runtime_diffusers-0.0.3.dist-info/METADATA,sha256=z3ZSTCX1eOifkl4AIZP32MO2YyYgaRTeZSqbiQjLSEY,841
+ai_simple_engine_runtime_diffusers-0.0.3.dist-info/WHEEL,sha256=M5asmiAlL6HEcOq52Yi5mmk9KmTVjY2RDPtO4p9DMrc,88
+ai_simple_engine_runtime_diffusers-0.0.3.dist-info/RECORD,,